Update app.py

app.py

@@ -1,59 +1,96 @@
+import os
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
+from diffusers import StableDiffusionPipeline
+from PIL import Image
 
-# Set Hugging Face token if using private models or rate limits apply
-# from huggingface_hub import login
-# login(token="your_huggingface_token")
-
-# Load Tamil to English translation model (use M2M100 for better support)
-from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
-translator_tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
-translator_model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
+# Load translation model/tokenizer (Tamil→English)
+try:
+    translator = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
+    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
+    tokenizer.src_lang = "ta"
+except Exception as e:
+    print(f"Error loading M2M100 model: {e}")
+    translator = tokenizer = None
 
-# Load text generation model
-text_generator = pipeline("text-generation", model="gpt2")
+# Load GPT-2 text generation pipeline
+try:
+    text_generator = pipeline("text-generation", model="gpt2")
+except Exception as e:
+    print(f"Error loading GPT-2 model: {e}")
+    text_generator = None
 
-# Load image generation model (e.g., SD 1.5)
-from diffusers import StableDiffusionPipeline
-import torch
+# Load Stable Diffusion pipeline
+hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", 
+        torch_dtype=torch.float16 if device=="cuda" else torch.float32,
+        use_auth_token=hf_token
+    )
+    pipe = pipe.to(device)
+    # Optionally enable efficient attention slicing if on GPU to save memory
+    if device == "cuda":
+        pipe.enable_attention_slicing()
+except Exception as e:
+    print(f"Error loading Stable Diffusion pipeline: {e}")
+    pipe = None
 
-image_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", 
-    torch_dtype=torch.float16
-)
-image_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+def tamil_to_image(tamil_text):
+    """
+    Translate Tamil text to English, generate new text with GPT-2, 
+    and produce an image with Stable Diffusion.
+    Returns (PIL.Image, info_text).
+    """
+    if not tamil_text or not tamil_text.strip():
+        return None, "Error: Please enter Tamil text as input."
 
-def process_input(tamil_text):
+    # Translation
     try:
-        # Step 1: Translate Tamil to English
-        translator_tokenizer.src_lang = "ta"
-        encoded = translator_tokenizer(tamil_text, return_tensors="pt")
-        generated_tokens = translator_model.generate(**encoded, forced_bos_token_id=translator_tokenizer.get_lang_id("en"))
-        english_text = translator_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-
-        # Step 2: Generate additional text
-        generated_output = text_generator(english_text, max_length=50, num_return_sequences=1)[0]["generated_text"]
+        tokenizer.src_lang = "ta"
+        encoded = tokenizer(tamil_text, return_tensors="pt")
+        generated_tokens = translator.generate(
+            **encoded, forced_bos_token_id=tokenizer.get_lang_id("en")
+        )
+        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    except Exception as e:
+        return None, f"Translation error: {e}"
 
-        # Step 3: Generate image from the final English text
-        image = image_pipe(generated_output).images[0]
+    # Text generation with GPT-2
+    try:
+        gen = text_generator(translation, max_length=50, num_return_sequences=1)
+        gen_text = gen[0]['generated_text'] if isinstance(gen, list) else gen['generated_text']
+    except Exception as e:
+        return None, f"Text generation error: {e}"
 
-        return english_text, generated_output, image
+    # Image generation with Stable Diffusion
+    try:
+        # Use the generated text as prompt
+        prompt = gen_text
+        if device == "cuda":
+            image = pipe(prompt, num_inference_steps=50).images[0]
+        else:
+            # On CPU, reduce steps to speed up if needed
+            image = pipe(prompt, num_inference_steps=25).images[0]
     except Exception as e:
-        return str(e), "", None
+        return None, f"Image generation error: {e}"
+
+    info = f"Translated → English: {translation}\nGPT-2 Prompt: {prompt}"
+    return image, info
 
-# Gradio interface
-demo = gr.Interface(
-    fn=process_input,
-    inputs=gr.Textbox(label="Enter Tamil Text"),
+# Build Gradio interface
+iface = gr.Interface(
+    fn=tamil_to_image,
+    inputs=gr.Textbox(label="Tamil Input", placeholder="Enter Tamil text here", type="text"),
     outputs=[
-        gr.Textbox(label="Translated English Text"),
-        gr.Textbox(label="Generated Description"),
-        gr.Image(label="Generated Image")
+        gr.Image(type="pil", label="Generated Image"),
+        gr.Textbox(label="Output Info", type="text")
     ],
-    title="Tamil to English → Text → Image Generator",
-    description="This app takes Tamil input, translates it to English, generates detailed text, and creates an image."
+    title="Tamil Text-to-Image Generator",
+    description="Enter Tamil text; this demo translates it to English, generates a story prompt with GPT-2, then creates an image with Stable Diffusion."
 )
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the app (in Spaces this will run on startup)
+iface.launch()