Update app.py

app.py

@@ -1,59 +1,59 @@
-from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
 import gradio as gr
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# Load the multilingual translation model (supports Tamil <-> English)
-model_name = "facebook/m2m100_418M"
-tokenizer = M2M100Tokenizer.from_pretrained(model_name)
-model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+# Set Hugging Face token if using private models or rate limits apply
+# from huggingface_hub import login
+# login(token="your_huggingface_token")
 
-# Load text generation pipeline (you can replace "gpt2" with your preferred model)
+# 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 text generation model
 text_generator = pipeline("text-generation", model="gpt2")
 
-def translate_tamil_to_english(text):
-    if not text.strip():
-        return ""
-    tokenizer.src_lang = "ta"
-    encoded = tokenizer(text, return_tensors="pt")
-    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id("en"))
-    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+# Load image generation model (e.g., SD 1.5)
+from diffusers import StableDiffusionPipeline
+import torch
 
-def process_text(tamil_text):
-    if not tamil_text.strip():
-        return "Please enter Tamil text.", 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 process_input(tamil_text):
     try:
         # Step 1: Translate Tamil to English
-        translation = translate_tamil_to_english(tamil_text)
-    except Exception as e:
-        return f"Translation error: {str(e)}", None, ""
+        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]
 
-    try:
-        # Step 2: Generate English text based on translation
-        generated = text_generator(translation, max_length=50, num_return_sequences=1)
-        generated_text = generated[0]['generated_text']
-    except Exception as e:
-        return translation, None, f"Text generation error: {str(e)}"
-
-    # Step 3: Placeholder for image generation
-    # Replace with your own image generation model/function
-    image_url = "https://via.placeholder.com/512.png?text=Generated+Image"
+        # Step 2: Generate additional text
+        generated_output = text_generator(english_text, max_length=50, num_return_sequences=1)[0]["generated_text"]
 
-    # Step 4: Description based on generated text (you can use an image captioning model here)
-    description = f"Generated description: {generated_text}"
+        # Step 3: Generate image from the final English text
+        image = image_pipe(generated_output).images[0]
 
-    return translation, image_url, description
+        return english_text, generated_output, image
+    except Exception as e:
+        return str(e), "", None
 
-iface = gr.Interface(
-    fn=process_text,
-    inputs=gr.Textbox(lines=2, label="Enter Tamil text"),
+# Gradio interface
+demo = gr.Interface(
+    fn=process_input,
+    inputs=gr.Textbox(label="Enter Tamil Text"),
     outputs=[
         gr.Textbox(label="Translated English Text"),
-        gr.Image(label="Generated Image"),
-        gr.Textbox(label="Image Description"),
+        gr.Textbox(label="Generated Description"),
+        gr.Image(label="Generated Image")
     ],
-    title="Tamil to English Translation + Text & Image Generation",
-    description="Enter Tamil text, get English translation, generated text, image, and description."
+    title="Tamil to English → Text → Image Generator",
+    description="This app takes Tamil input, translates it to English, generates detailed text, and creates an image."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()