Update app.py

app.py

@@ -33,13 +33,13 @@ pipe = FluxWithCFGPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", t
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
 pipe.to(device)
 
-# --- Load Multilingual Translator (M2M100) ---
+# --- Load M2M100 Albanian-English Translator ---
 try:
     tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
     model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M").to(device)
     print("✅ Loaded M2M100 Albanian-English translator")
 except Exception as e:
-    print(f"❌ Failed to load M2M100: {e}")
+    print(f"❌ Failed to load M2M100 translator: {e}")
     tokenizer = None
     model = None
 
@@ -48,15 +48,15 @@ def translate_sq_to_en(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]
+        encoded = tokenizer(text, return_tensors="pt", padding=True).to(device)
+        generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id("en"))
+        translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
         return translated
     except Exception as e:
         print(f"❌ Translation failed: {e}")
         return text
 
-# --- Inference Function ---
+# --- Main 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:
@@ -71,8 +71,10 @@ 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))
 
     prompt_final = prompt.strip()
-    if any(c in prompt_final for c in "ëçËÇ") or prompt_final.lower().startswith("një ") or prompt_final.lower().count(" ") >= 2:
-        print("🌐 Detected likely Albanian. Translating...")
+
+    # Detect if prompt is probably Albanian by common Albanian letters or phrase starts
+    if any(c in prompt_final for c in "ëçËÇ") or prompt_final.lower().startswith("një "):
+        print(f"🌐 Detected likely Albanian. Translating prompt: {prompt_final}")
         prompt_final = translate_sq_to_en(prompt_final)
         print(f"✅ Translated prompt: {prompt_final}")
     else: