Update app.py

app.py

@@ -16,7 +16,6 @@ torch.backends.cudnn.benchmark = True
 
 # --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
 DEFAULT_WIDTH = 1024
 DEFAULT_HEIGHT = 576
 ASPECT_RATIOS = {
@@ -30,15 +29,18 @@ INFERENCE_STEPS = 8
 dtype = torch.float16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+print("⏳ Loading Flux pipeline...")
 pipe = FluxWithCFGPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype)
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
 pipe.to(device)
+print("✅ Flux pipeline loaded.")
 
-# --- Load Multilingual Translator (M2M100) ---
+# --- Load M2M100 Translator ---
 try:
+    print("⏳ Loading M2M100 tokenizer and model...")
     tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
     model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M").to(device)
-    print("✅ Loaded M2M100 Albanian-English translator")
+    print("✅ M2M100 loaded.")
 except Exception as e:
     print(f"❌ Failed to load M2M100: {e}")
     tokenizer = None
@@ -46,12 +48,14 @@ except Exception as e:
 
 def translate_sq_to_en(text):
     if not tokenizer or not model:
+        print("⚠️ Translator not loaded, returning original text")
         return text
     try:
         tokenizer.src_lang = "sq"
         encoded = tokenizer(text, return_tensors="pt").to(device)
         generated = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id("en"))
         translated = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
+        print(f"🌐 Translation successful: {translated}")
         return translated
     except Exception as e:
         print(f"❌ Translation failed: {e}")
@@ -60,22 +64,22 @@ def translate_sq_to_en(text):
 def is_albanian(text):
     try:
         lang = detect(text)
+        print(f"🕵️ Language detected: {lang}")
         return lang == "sq"
     except Exception as e:
         print(f"⚠️ Language detection failed: {e}")
         return False
 
-# --- Inference Function ---
 @spaces.GPU
 def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", randomize_seed: bool = False):
     if pipe is None:
         raise gr.Error("Pipeline nuk u ngarkua.")
 
-    if not prompt or prompt.strip() == "":
+    prompt_clean = prompt.strip()
+    if not prompt_clean:
         return None, seed, "Gabim: Plotësoni përshkrimin."
 
-    # Enforce Albanian only input
-    if not is_albanian(prompt.strip()):
+    if not is_albanian(prompt_clean):
         return None, seed, "Ju lutemi shkruani vetëm në gjuhën shqipe."
 
     if randomize_seed:
@@ -84,16 +88,18 @@ def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", rand
     width, height = ASPECT_RATIOS.get(aspect_ratio, (DEFAULT_WIDTH, DEFAULT_HEIGHT))
 
     # Translate Albanian prompt to English
-    prompt_final = translate_sq_to_en(prompt.strip())
-    print(f"🌐 Translated prompt: {prompt_final}")
+    prompt_final = translate_sq_to_en(prompt_clean)
 
+    # Add quality tags for generation
     prompt_final += ", ultra realistic, sharp, 8k resolution"
 
+    print(f"🎯 Final prompt for generation: {prompt_final}")
+
     try:
-        generator = torch.Generator(device=device).manual_seed(int(float(seed)))
+        generator = torch.Generator(device=device).manual_seed(int(seed))
         start_time = time.time()
         with torch.inference_mode():
-            image = pipe(
+            images = pipe(
                 prompt=prompt_final,
                 width=width,
                 height=height,
@@ -101,7 +107,8 @@ def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", rand
                 generator=generator,
                 output_type="pil",
                 return_dict=False
-            )[0][0]
+            )
+        image = images[0][0]
         latency = time.time() - start_time
         status = f"Koha e përpunimit: {latency:.2f} sekonda"
         return image, seed, status
@@ -140,7 +147,7 @@ button[aria-label="Download"] {
         with gr.Column(scale=2):
             output_image = gr.Image(label="Imazhi i Gjeneruar", interactive=False, show_download_button=True)
         with gr.Column(scale=1):
-            prompt = gr.Text(label="Përshkrimi (vetëm në shqip)", placeholder="Shkruani vetëm në gjuhën shqipe...", lines=3)
+            prompt_input = gr.Text(label="Përshkrimi (vetëm në shqip)", placeholder="Shkruani vetëm në gjuhën shqipe...", lines=3)
             generate_btn = gr.Button("🎨 Gjenero")
             aspect_ratio = gr.Radio(label="Raporti i Imazhit", choices=list(ASPECT_RATIOS.keys()), value="16:9")
             randomize_seed = gr.Checkbox(label="Përdor numër të rastësishëm", value=True)
@@ -150,7 +157,7 @@ button[aria-label="Download"] {
     gr.Examples(
         examples=examples,
         fn=generate_image,
-        inputs=[prompt],
+        inputs=[prompt_input],
         outputs=[output_image, gr.Number(visible=False), latency],
         cache_examples=True,
         cache_mode="eager"
@@ -158,7 +165,7 @@ button[aria-label="Download"] {
 
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, gr.Number(value=42, visible=False), aspect_ratio, randomize_seed],
+        inputs=[prompt_input, gr.Number(value=42, visible=False), aspect_ratio, randomize_seed],
         outputs=[output_image, gr.Number(visible=False), latency],
         show_progress="full"
     )