pushing experiment code

app.py

@@ -26,34 +26,144 @@
 #     interface.launch()
 
 
+# import gradio as gr
+# from translate import translate_text
+# from image_generator import generate_image
+
+# def chatbot(input_text, src_lang="auto"):
+#     """Process user input, translate it, and generate an image."""
+#     # Translate input to English
+#     translated_text = translate_text(input_text, src_lang, "en")
+    
+#     # Generate an image based on the translated text
+#     image_path = generate_image(f"A scene depicting: {translated_text}")
+    
+#     return translated_text, image_path
+
+# # Gradio Interface
+# interface = gr.Interface(
+#     fn=chatbot,
+#     inputs=[
+#         gr.Textbox(label="Enter text in any language"),
+#         gr.Textbox(label="Source Language (optional, e.g., 'es' for Spanish)", value="auto"),
+#     ],
+#     outputs=[
+#         gr.Textbox(label="Translated Text"),
+#         gr.Image(label="Generated Image"),
+#     ],
+#     title="LinguaVision - Multilingual Chatbot with Image Generation",
+#     description="Enter text in any language to translate it into English and generate an image based on the text.\n NOTE: This tool takes approximately 12 minutes to execute.",
+# )
+
+# if __name__ == "__main__":
+#     interface.launch(share=True)
+
+import torch
+from transformers import MarianMTModel, MarianTokenizer
+from diffusers import StableDiffusionPipeline, DDIMScheduler
 import gradio as gr
-from translate import translate_text
-from image_generator import generate_image
+from typing import Tuple, Optional
+import logging
+
+class LinguaVisionSystem:
+    def __init__(self, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
+        self.device = device
+        self.logger = logging.getLogger(__name__)
+        
+        # Initialize translation pipeline
+        self.translation_config = {
+            "model_name": "Helsinki-NLP/opus-mt-mul-en",
+            "max_length": 128,
+            "num_beams": 4
+        }
+        self._init_translation_pipeline()
+        
+        # Initialize image generation pipeline
+        self.image_config = {
+            "model_id": "stabilityai/stable-diffusion-2-1-base",
+            "safety_checker": None,  # Disable for performance
+            "scheduler": DDIMScheduler
+        }
+        self._init_image_pipeline()
+
+    def _init_translation_pipeline(self) -> None:
+        try:
+            self.tokenizer = MarianTokenizer.from_pretrained(
+                self.translation_config["model_name"]
+            )
+            self.translation_model = MarianMTModel.from_pretrained(
+                self.translation_config["model_name"]
+            ).to(self.device)
+        except Exception as e:
+            self.logger.error(f"Translation pipeline initialization failed: {e}")
+            raise
+
+    def _init_image_pipeline(self) -> None:
+        try:
+            self.image_pipeline = StableDiffusionPipeline.from_pretrained(
+                self.image_config["model_id"],
+                scheduler=self.image_config["scheduler"](),
+                safety_checker=self.image_config["safety_checker"]
+            ).to(self.device)
+        except Exception as e:
+            self.logger.error(f"Image pipeline initialization failed: {e}")
+            raise
+
+    @torch.inference_mode()
+    def translate_text(self, text: str) -> Optional[str]:
+        try:
+            inputs = self.tokenizer(
+                text, 
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.translation_config["max_length"]
+            ).to(self.device)
+            
+            translated = self.translation_model.generate(
+                **inputs,
+                num_beams=self.translation_config["num_beams"],
+                early_stopping=True
+            )
+            
+            return self.tokenizer.decode(translated[0], skip_special_tokens=True)
+        except Exception as e:
+            self.logger.error(f"Translation failed: {e}")
+            return None
+
+    def process_input(self, text: str) -> Tuple[str, str]:
+        translated_text = self.translate_text(text)
+        if translated_text:
+            image = self.image_pipeline(
+                prompt=f"A photorealistic scene depicting: {translated_text}",
+                num_inference_steps=50,
+                guidance_scale=7.5
+            ).images[0]
+            
+            image_path = "output.png"
+            image.save(image_path)
+            return translated_text, image_path
+        return "Translation failed", None
 
-def chatbot(input_text, src_lang="auto"):
-    """Process user input, translate it, and generate an image."""
-    # Translate input to English
-    translated_text = translate_text(input_text, src_lang, "en")
+def create_interface() -> gr.Interface:
+    system = LinguaVisionSystem()
     
-    # Generate an image based on the translated text
-    image_path = generate_image(f"A scene depicting: {translated_text}")
+    interface = gr.Interface(
+        fn=system.process_input,
+        inputs=gr.Textbox(
+            label="Enter text in any language",
+            placeholder="Type your text here..."
+        ),
+        outputs=[
+            gr.Textbox(label="English Translation"),
+            gr.Image(label="Generated Visualization")
+        ],
+        title="LinguaVision: AI-Powered Language Learning Assistant",
+        description="Transform text into visuals for enhanced language learning"
+    )
     
-    return translated_text, image_path
-
-# Gradio Interface
-interface = gr.Interface(
-    fn=chatbot,
-    inputs=[
-        gr.Textbox(label="Enter text in any language"),
-        gr.Textbox(label="Source Language (optional, e.g., 'es' for Spanish)", value="auto"),
-    ],
-    outputs=[
-        gr.Textbox(label="Translated Text"),
-        gr.Image(label="Generated Image"),
-    ],
-    title="LinguaVision - Multilingual Chatbot with Image Generation",
-    description="Enter text in any language to translate it into English and generate an image based on the text.\n NOTE: This tool takes approximately 12 minutes to execute.",
-)
+    return interface
 
 if __name__ == "__main__":
-    interface.launch(share=True)
+    interface = create_interface()
+    interface.launch(server_name="0.0.0.0", server_port=7860)

image_generator.py

@@ -1,50 +0,0 @@
-# from diffusers import StableDiffusionPipeline
-
-# def generate_image(prompt):
-#     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-#     model.to("cuda")  # Use GPU for faster generation
-#     image = model(prompt).images[0]
-#     image.save("output.png")
-#     return "output.png"
-
-# if __name__ == "__main__":
-#     prompt = "A friendly person saying 'How are you?'"
-#     print("Generated Image Path:", generate_image(prompt))
-
-
-# from diffusers import StableDiffusionPipeline
-# import torch
-
-# def generate_image(prompt):
-#     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    
-#     # Use GPU if available, otherwise fallback to CPU
-#     device = "cuda" if torch.cuda.is_available() else "cpu"
-#     model.to(device)
-    
-#     image = model(prompt).images[0]
-#     image.save("output.png")
-#     return "output.png"
-
-# if __name__ == "__main__":
-#     prompt = "A friendly person saying 'How are you?'"
-#     print("Generated Image Path:", generate_image(prompt))
-
-
-from diffusers import StableDiffusionPipeline
-import torch
-
-# Preload the model globally
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1-base",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-)
-model.to(device)
-
-def generate_image(prompt):
-    """Generate an image from a text prompt."""
-    image = model(prompt).images[0]
-    output_path = "output.png"
-    image.save(output_path)
-    return output_path

translate.py

@@ -1,27 +0,0 @@
-# from transformers import MarianMTModel, MarianTokenizer
-
-# def translate_text(text, src_lang="es", tgt_lang="en"):
-#     model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
-#     tokenizer = MarianTokenizer.from_pretrained(model_name)
-#     model = MarianMTModel.from_pretrained(model_name)
-#     inputs = tokenizer(text, return_tensors="pt", padding=True)
-#     translated = model.generate(**inputs)
-#     return tokenizer.decode(translated[0], skip_special_tokens=True)
-
-# if __name__ == "__main__":
-#     input_text = "¿Cómo estás?"
-#     print("Translated Text:", translate_text(input_text, src_lang="es", tgt_lang="en"))
-
-
-from transformers import MarianMTModel, MarianTokenizer
-
-# Preload the translation model globally
-model_name = "Helsinki-NLP/opus-mt-mul-en"
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-translation_model = MarianMTModel.from_pretrained(model_name)
-
-def translate_text(text, src_lang="auto", tgt_lang="en"):
-    """Translate text from any language to English."""
-    inputs = tokenizer(text, return_tensors="pt", padding=True)
-    translated = translation_model.generate(**inputs)
-    return tokenizer.decode(translated[0], skip_special_tokens=True)