Update app.py

app.py

@@ -6,6 +6,7 @@ import torch
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
 from custom_pipeline import FluxWithCFGPipeline
+from transformers import pipeline  # Added for translation
 
 # --- Torch Optimizations ---
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -31,6 +32,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)
 
+# --- Translation Pipeline ---
+try:
+    translator = pipeline("translation", model="Helsinki-NLP/opus-mt-sq-en", force_download=True)
+except Exception as e:
+    print(f"Warning: Failed to load translation model: {e}. Using input prompt directly.")
+    translator = None
+
 # --- Inference Function ---
 @spaces.GPU
 def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", randomize_seed: bool = False):
@@ -45,14 +53,23 @@ 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 += ", ultra realistic, sharp, 8k resolution, cinematic lighting"
+    # Translate Albanian prompt to English if available
+    prompt_final = prompt.strip()
+    if translator and prompt_final:
+        try:
+            prompt_final = translator(prompt_final)[0]['translation_text']
+        except Exception as e:
+            print(f"Warning: Translation failed: {e}. Using input prompt directly.")
+            prompt_final = prompt.strip()
+    
+    prompt_final += ", ultra realistic, sharp, 8k resolution, cinematic lighting"
 
     try:
         generator = torch.Generator(device=device).manual_seed(int(float(seed)))
         start_time = time.time()
         with torch.inference_mode():
             image = pipe(
-                prompt=prompt,
+                prompt=prompt_final,
                 width=width,
                 height=height,
                 num_inference_steps=INFERENCE_STEPS,
@@ -74,8 +91,9 @@ examples = [
     "Pamje nga një shtëpi moderne në stilin Minecraft"
 ]
 
+# --- App Layout ---
 with gr.Blocks(css="""
-body::before {
+body::before, body::after {
     content: "";
     display: block;
     height: 640px;
@@ -95,7 +113,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", placeholder="Shkruani se çfarë doni të krijoni...", lines=3)
+            prompt = gr.Text(label="Përshkrimi (Shqip ose Anglisht)", placeholder="Shkruani se çfarë doni të krijoni...", 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)
@@ -118,4 +136,4 @@ button[aria-label="Download"] {
         show_progress="full"
     )
 
-app.launch(share=True)
+app.launch(share=True)