Update app.py

app.py

@@ -3,26 +3,29 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from diffusers import StableDiffusionPipeline
 from sentence_transformers import SentenceTransformer, util
 import torch
+import contextlib
 
 # --- Load models ---
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Text-to-text model
-text_model_name = "google/flan-t5-small"
+text_model_name = "google/flan-t5-large"
 text_tokenizer = AutoTokenizer.from_pretrained(text_model_name)
-text_model = AutoModelForSeq2SeqLM.from_pretrained(text_model_name)
+text_model = AutoModelForSeq2SeqLM.from_pretrained(text_model_name).to(device)
 
 # Text-to-image model
 image_model_id = "runwayml/stable-diffusion-v1-5"
 image_pipe = StableDiffusionPipeline.from_pretrained(
-    image_model_id, torch_dtype=torch.float16 if device=="cuda" else torch.float32
+    image_model_id,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    safety_checker=None  # Optional for debugging
 )
 image_pipe = image_pipe.to(device)
 
 # Sentence similarity model
 embedder = SentenceTransformer('all-MiniLM-L6-v2')
 
-# Example prompts that mean “generate image”
+# Image-like trigger phrases
 image_triggers = [
     "generate an image of",
     "draw a",
@@ -33,30 +36,30 @@ image_triggers = [
     "sketch",
 ]
 
-# --- Main Logic ---
+# --- Core logic ---
 def multimodal_agent(prompt):
-    # Step 1: Check similarity to image prompts
+    # Step 1: Semantic similarity to image triggers
     prompt_embedding = embedder.encode(prompt, convert_to_tensor=True)
     trigger_embeddings = embedder.encode(image_triggers, convert_to_tensor=True)
-    
     cosine_scores = util.pytorch_cos_sim(prompt_embedding, trigger_embeddings)
     max_score = torch.max(cosine_scores).item()
 
-    # Step 2: Threshold to decide
+    # Step 2: Decision branch
     if max_score > 0.65:
         # Generate image
-        image = image_pipe(prompt).images[0]
-        return None, image  # Return image only
+        with torch.autocast("cuda") if device == "cuda" else contextlib.nullcontext():
+            image = image_pipe(prompt).images[0]
+        return None, image
     else:
         # Generate text
-        inputs = text_tokenizer(prompt, return_tensors="pt")
+        inputs = text_tokenizer(prompt, return_tensors="pt").to(device)
         outputs = text_model.generate(**inputs, max_new_tokens=100)
         text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return text, None  # Return text only
+        return text, None
 
-# --- Gradio UI ---
+# --- UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Smart Multimodal AI Agent\nDecides text or image based on your prompt.")
+    gr.Markdown("# 🤖 Smart Multimodal AI Agent\nGive a prompt — It decides text vs image automatically!")
 
     input_box = gr.Textbox(label="Enter your prompt")
     output_text = gr.Textbox(label="Text Output")