Create app.py

app.py

@@ -0,0 +1,51 @@
+"""
+app.py
+"""
+
+import gradio as gr
+import spaces
+import torch
+from transformers import (
+    LlavaNextProcessor,
+    LlavaNextForConditionalGeneration
+)
+
+# Load model and processor
+model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-34b-hf",
+                                                          torch_dtype=torch.float16,
+                                                          low_cpu_mem_usage=True)
+model.to("cuda:0")
+processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-34b-hf")
+
+@spaces.GPU(duration=120)
+def generate_caption(image):
+    """
+    Generate a poem from an image.
+    """
+    # Process the image and the prompt
+    inputs = processor(texts=["Write a poem about this picture"],
+                       images=[image],
+                       return_tensors="pt").to('cuda')
+    # Generate the output
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=128,
+        )
+    prompt_len = inputs["input_ids"].shape[1]
+    decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
+    return decoded_text
+
+
+# Define the Gradio interface
+description = """This is a demo of [`llava-hf/llava-v1.6-34b-hf`](https://huggingface.co/llava-hf/llava-v1.6-34b-hf) hosted with ZeroGPU."""
+
+iface = gr.Interface(
+    fn=generate_caption,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Textbox(label="Generated Poem"),
+    description=description
+)
+
+# Launch the interface
+iface.launch()