Update app.py

app.py

@@ -4,48 +4,65 @@ from diffusers import StableDiffusionPipeline
 import torch
 import os
 
-# ✅ Load HF token safely from environment
+# 1. Use Hugging Face token securely
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
-# ✅ Select device
+# 2. Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ✅ 1. Tamil → English Translator (multilingual model)
-translator = pipeline(
-    "translation",
-    model="Helsinki-NLP/opus-mt-mul-en",
-    use_auth_token=HF_TOKEN
-)
+# 3. Load translator (Tamil → English using multilingual model)
+try:
+    translator = pipeline(
+        "translation",
+        model="Helsinki-NLP/opus-mt-mul-en",
+        use_auth_token=HF_TOKEN
+    )
+except Exception as e:
+    translator = None
+    print(f"Error loading translator: {e}")
 
-# ✅ 2. English Text Generator
-generator = pipeline("text-generation", model="gpt2")
+# 4. Load GPT2 for English text generation
+try:
+    generator = pipeline("text-generation", model="gpt2")
+except Exception as e:
+    generator = None
+    print(f"Error loading GPT2: {e}")
 
-# ✅ 3. Image Generator using Stable Diffusion
-image_pipe = StableDiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4",
-    use_auth_token=HF_TOKEN,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-)
-image_pipe = image_pipe.to(device)
+# 5. Load Stable Diffusion for image generation
+try:
+    image_pipe = StableDiffusionPipeline.from_pretrained(
+        "CompVis/stable-diffusion-v1-4",
+        use_auth_token=HF_TOKEN,
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    ).to(device)
+except Exception as e:
+    image_pipe = None
+    print(f"Error loading Stable Diffusion: {e}")
+
+# 6. Full pipeline function with safe error handling
+def generate_image_from_tamil(tamil_text):
+    if not translator or not generator or not image_pipe:
+        return "Model load error", "Model load error", None
 
-# ✅ Function to connect all 3 stages
-def generate_image_from_tamil(tamil_input):
-    translated = translator(tamil_input, max_length=100)[0]['translation_text']
-    generated = generator(translated, max_length=50, num_return_sequences=1)[0]['generated_text'].strip()
-    image = image_pipe(generated).images[0]
-    return translated, generated, image
+    try:
+        translated = translator(tamil_text, max_length=100)[0]['translation_text']
+        prompt = generator(translated, max_length=50, num_return_sequences=1)[0]['generated_text']
+        image = image_pipe(prompt).images[0]
+        return translated, prompt, image
+    except Exception as e:
+        return f"Translation/Image generation error: {str(e)}", "", None
 
-# ✅ Gradio Interface
+# 7. Gradio UI
 iface = gr.Interface(
     fn=generate_image_from_tamil,
     inputs=gr.Textbox(lines=2, label="Enter Tamil Text"),
     outputs=[
         gr.Textbox(label="Translated English Text"),
-        gr.Textbox(label="Generated English Prompt"),
+        gr.Textbox(label="Generated Prompt"),
         gr.Image(label="Generated Image")
     ],
     title="Tamil to Image Generator",
-    description="🔤 Translates Tamil → English → Generates prompt → Creates an image using Stable Diffusion"
+    description="Translate Tamil ➜ Generate English Text ➜ Create Image"
 )
 
 iface.launch()