Create app.py

app.py

@@ -0,0 +1,68 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from diffusers import StableDiffusionPipeline
+from sentence_transformers import SentenceTransformer, util
+import torch
+
+# --- Load models ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Text-to-text model
+text_model_name = "google/flan-t5-small"
+text_tokenizer = AutoTokenizer.from_pretrained(text_model_name)
+text_model = AutoModelForSeq2SeqLM.from_pretrained(text_model_name)
+
+# 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_pipe = image_pipe.to(device)
+
+# Sentence similarity model
+embedder = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Example prompts that mean “generate image”
+image_triggers = [
+    "generate an image of",
+    "draw a",
+    "create a picture of",
+    "show me a",
+    "visualize",
+    "render",
+    "sketch",
+]
+
+# --- Main Logic ---
+def multimodal_agent(prompt):
+    # Step 1: Check similarity to image prompts
+    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
+    if max_score > 0.65:
+        # Generate image
+        image = image_pipe(prompt).images[0]
+        return None, image  # Return image only
+    else:
+        # Generate text
+        inputs = text_tokenizer(prompt, return_tensors="pt")
+        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
+
+# --- Gradio UI ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Smart Multimodal AI Agent\nDecides text or image based on your prompt.")
+
+    input_box = gr.Textbox(label="Enter your prompt")
+    output_text = gr.Textbox(label="Text Output")
+    output_image = gr.Image(label="Image Output")
+
+    btn = gr.Button("Generate")
+    btn.click(multimodal_agent, inputs=input_box, outputs=[output_text, output_image])
+
+demo.launch()