Create app.py

app.py

@@ -0,0 +1,83 @@
+import torch
+from PIL import Image
+import streamlit as st
+from transformers import (
+    CLIPProcessor, CLIPModel,
+    DistilBertTokenizer, DistilBertModel,
+    GPT2LMHeadModel, GPT2Tokenizer
+)
+
+# -------- Load Models --------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").vision_model.to(device)
+
+text_tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+text_encoder = DistilBertModel.from_pretrained("distilbert-base-uncased").to(device)
+
+decoder_tokenizer = GPT2Tokenizer.from_pretrained("sreebhargavibalija/sreebhargavibalija-multimodal-gen")
+decoder_tokenizer.pad_token = decoder_tokenizer.eos_token
+decoder = GPT2LMHeadModel.from_pretrained("sreebhargavibalija/sreebhargavibalija-multimodal-gen").to(device)
+
+# -------- Fusion Wrapper --------
+class MultimodalGenerator(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.image_encoder = clip_model
+        self.text_encoder = text_encoder
+        self.decoder = decoder
+
+        self.project_image = torch.nn.Linear(768, 768)
+        self.project_text = torch.nn.Linear(768, 768)
+        self.fusion = torch.nn.Linear(768 * 2, 768)
+
+    def forward(self, image_tensor, prompt_input_ids, prompt_attention_mask, max_length=50):
+        img_feat = self.image_encoder(pixel_values=image_tensor).last_hidden_state[:, 0, :]
+        img_feat = self.project_image(img_feat)
+
+        txt_feat = self.text_encoder(input_ids=prompt_input_ids, attention_mask=prompt_attention_mask).last_hidden_state[:, 0, :]
+        txt_feat = self.project_text(txt_feat)
+
+        fused = self.fusion(torch.cat([img_feat, txt_feat], dim=-1)).unsqueeze(1)
+
+        generated = self.decoder.generate(
+            inputs_embeds=fused,
+            max_length=max_length,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            num_return_sequences=1,
+            pad_token_id=self.decoder.config.pad_token_id
+        )
+        return generated
+
+# Initialize model
+model = MultimodalGenerator().to(device)
+model.eval()
+
+# -------- Streamlit UI --------
+st.set_page_config(page_title="Multimodal LLM", layout="centered")
+st.title("🧠 Multimodal LLM: Image + Prompt → Text")
+
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+prompt_text = st.text_input("Enter your prompt (e.g. 'Describe this scene'):")
+
+if uploaded_file is not None and prompt_text.strip():
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    image_tensor = clip_processor(images=image, return_tensors="pt")["pixel_values"].to(device)
+    prompt_inputs = text_tokenizer(prompt_text, return_tensors="pt", padding="max_length", truncation=True, max_length=64)
+    prompt_ids = prompt_inputs["input_ids"].to(device)
+    prompt_mask = prompt_inputs["attention_mask"].to(device)
+
+    with st.spinner("Generating..."):
+        with torch.no_grad():
+            generated_ids = model(image_tensor, prompt_ids, prompt_mask, max_length=64)
+            output_text = decoder_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+    st.markdown("### ✨ Generated Text")
+    st.success(output_text)
+else:
+    st.info("👆 Upload an image and enter a prompt to get started!")