Update app.py

app.py

@@ -8,6 +8,7 @@ import os
 from diffusers import DiffusionPipeline, AutoencoderTiny
 from custom_pipeline import FluxWithCFGPipeline
 from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+from langdetect import detect
 
 # --- Torch Optimizations ---
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -33,13 +34,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 M2M100 Albanian-English Translator ---
+# --- Load Multilingual Translator (M2M100) ---
 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 translator: {e}")
+    print(f"❌ Failed to load M2M100: {e}")
     tokenizer = None
     model = None
 
@@ -48,37 +49,43 @@ def translate_sq_to_en(text):
         return text
     try:
         tokenizer.src_lang = "sq"
-        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]
+        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]
         return translated
     except Exception as e:
         print(f"❌ Translation failed: {e}")
         return text
 
-# --- Main Inference Function ---
+def is_albanian(text):
+    try:
+        lang = detect(text)
+        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("Pipelinei nuk u ngarkua.")
+        raise gr.Error("Pipeline nuk u ngarkua.")
 
     if not prompt or prompt.strip() == "":
         return None, seed, "Gabim: Plotësoni përshkrimin."
 
+    # Enforce Albanian only input
+    if not is_albanian(prompt.strip()):
+        return None, seed, "Ju lutemi shkruani vetëm në gjuhën shqipe."
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     width, height = ASPECT_RATIOS.get(aspect_ratio, (DEFAULT_WIDTH, DEFAULT_HEIGHT))
 
-    prompt_final = prompt.strip()
-
-    # 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:
-        print("🟢 Prompt seems English. Skipping translation.")
+    # Translate Albanian prompt to English
+    prompt_final = translate_sq_to_en(prompt.strip())
+    print(f"🌐 Translated prompt: {prompt_final}")
 
     prompt_final += ", ultra realistic, sharp, 8k resolution"
 
@@ -101,7 +108,7 @@ def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", rand
     except Exception as e:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-        raise gr.Error(f"❌ Gabim gjatë gjenerimit: {e}")
+        raise gr.Error(f"Gabim gjatë gjenerimit: {e}")
 
 # --- Examples ---
 examples = [
@@ -133,7 +140,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 (Shqip ose Anglisht)", placeholder="Shkruani se çfarë doni të krijoni...", lines=3)
+            prompt = 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)