Update app.py

app.py

@@ -20,32 +20,30 @@ import spaces  # Hugging Face Spaces module
 
 # Initialize the model
 pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")
-st.title("Image Caption Generator")
-st.write("Upload an image or provide an image URL to generate its caption.")
-
-# Option for image upload
-img_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
-if img_file is not None:
-    raw_image = Image.open(img_file).convert('RGB')
-text = "a photography of"
-inputs = processor(raw_image, text, return_tensors="pt",  padding =True, truncation=True, max_length =250)
-out = model.generate(**inputs)
-caption = processor.decode(out[0], skip_special_tokens=True)
+
+
 
 @spaces.GPU(duration=300)
-def generate_image(prompt):
-    # Move the model to GPU if available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe.to(device)
-    image = pipe(prompt).images[0]
-    return image
-
-# Create the Gradio interface
-iface = gr.Interface(fn=generate_image, 
-                     inputs=caption, 
-                     outputs=gr.Image(label="Generated Image"),
-                     title="Astronaut in a Jungle Model")
-
-# Launch the interface
-iface.launch(share=True)
+def generate_caption_and_image(image):
+    # Process the image
+    raw_image = image.convert("RGB")
+    
+    # Generate caption
+    inputs = processor(raw_image, return_tensors="pt", padding=True, truncation=True, max_length=250)
+    inputs = {key: val.to(device) for key, val in inputs.items()}
+    out = model.generate(**inputs)
+    caption = processor.decode(out[0], skip_special_tokens=True)
+
+    # Generate image based on the caption
+    generated_image = pipe(caption).images[0]
+
+    return caption, generated_image
+
+# Gradio UI
+iface = gr.Interface(
+    fn=generate_caption_and_image,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=[gr.Textbox(label="Generated Caption"), gr.Image(label="Generated Design")],
+    live=True
+)