Spaces:
Sleeping
Sleeping
File size: 2,412 Bytes
12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 27d6624 12d2cf1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import gradio as gr
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-large"
text_tokenizer = AutoTokenizer.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,
safety_checker=None # Optional for debugging
)
image_pipe = image_pipe.to(device)
# Sentence similarity model
embedder = SentenceTransformer('all-MiniLM-L6-v2')
# Image-like trigger phrases
image_triggers = [
"generate an image of",
"draw a",
"create a picture of",
"show me a",
"visualize",
"render",
"sketch",
]
# --- Core logic ---
def multimodal_agent(prompt):
# 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: Decision branch
if max_score > 0.65:
# Generate image
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").to(device)
outputs = text_model.generate(**inputs, max_new_tokens=100)
text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
return text, None
# --- UI ---
with gr.Blocks() as demo:
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")
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() |