Create app.py

app.py

@@ -0,0 +1,19 @@
+import torch
+from transformers import LlavaProcessor, LlavaForConditionalGeneration
+import gradio as gr
+from PIL import Image
+
+# Load LLaVA model and processor
+model_id = "llava-hf/llava-1.5-7b"
+processor = LlavaProcessor.from_pretrained(model_id)
+model = LlavaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+
+# Function to generate captions
+def generate_caption(image, prompt="Describe this image."):
+    inputs = processor(images=image, text=prompt, return_tensors="pt").to("cuda")
+    output = model.generate(**inputs, max_new_tokens=50)
+    return processor.batch_decode(output, skip_special_tokens=True)[0]
+
+# Gradio UI
+demo = gr.Interface(fn=generate_caption, inputs=[gr.Image(type="pil"), gr.Textbox()], outputs="text")
+demo.launch()