Update app.py

app.py

@@ -1,26 +1,45 @@
 import gradio as gr
-from transformers import CLIPProcessor, CLIPModel
+from transformers import CLIPProcessor, CLIPModel, AutoTokenizer, AutoModelForCausalLM
 from PIL import Image
 import torch
 
-# Load CLIP
+# Load CLIP model
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# Prompt template
+# Load GPT-2 (or any captioning LLM)
+lm_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+lm_model = AutoModelForCausalLM.from_pretrained("gpt2")
+
 def generate_caption(image):
-    inputs = clip_processor(images=image, return_tensors="pt")
-    outputs = clip_model.get_image_features(**inputs)
-    
-    # Convert image features into a dummy "caption" using top concept labels
-    # (In actual implementation, this could be passed to GPT-like models)
-    # Here we simulate a caption
-    return "A photo showing something relevant to the content."
-
-demo = gr.Interface(fn=generate_caption,
-                    inputs=gr.Image(type="pil"),
-                    outputs="text",
-                    title="Image Captioning with CLIP & GPT-style Generation",
-                    description="Upload an image to get a descriptive caption. Based on CLIP for vision understanding.")
-
-demo.launch()
+    if image is None:
+        return "No image uploaded."
+
+    # Create basic prompt ideas for CLIP to compare
+    concepts = ["cat", "dog", "person", "landscape", "food", "technology", "vehicle", "building", "nature"]
+    prompts = [f"an image of a {c}" for c in concepts]
+
+    # Use CLIP to find the best concept
+    inputs = clip_processor(text=prompts, images=image, return_tensors="pt", padding=True)
+    outputs = clip_model(**inputs)
+    probs = outputs.logits_per_image.softmax(dim=1)
+    best = torch.argmax(probs).item()
+    selected_concept = concepts[best]
+
+    # Use the concept as seed for GPT caption generation
+    gpt_prompt = f"This is an image of a {selected_concept}. It shows"
+    input_ids = lm_tokenizer.encode(gpt_prompt, return_tensors="pt")
+    gpt_output = lm_model.generate(input_ids, max_length=30, do_sample=True, top_k=50, top_p=0.95)
+    generated_text = lm_tokenizer.decode(gpt_output[0], skip_special_tokens=True)
+
+    return generated_text
+
+iface = gr.Interface(
+    fn=generate_caption,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(label="Generated Caption"),
+    title="Image Captioning with CLIP + GPT",
+    description="CLIP guesses image context, GPT generates free-text caption."
+)
+
+iface.launch()