Update app.py

app.py

@@ -1,8 +1,8 @@
 import os
 import requests
 from PIL import Image
-import streamlit as st
 import torch
+import gradio as gr
 from huggingface_hub import login
 from transformers import AutoProcessor, AutoModelForCausalLM
 from diffusers import DiffusionPipeline
@@ -21,44 +21,34 @@ caption_model_name = "pretrained-caption-model"  # Replace with the actual model
 processor = AutoProcessor.from_pretrained(caption_model_name)
 model = AutoModelForCausalLM.from_pretrained(caption_model_name)
 
-# Move models to GPU if available
+# Check for GPU availability (handled automatically by Hugging Face Spaces)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe.to(device)
 model.to(device)
 
-# Streamlit UI
-st.title("Image Caption and Design Generator")
-st.write("Upload an image or provide an image URL to generate a caption and use it to create a similar design.")
+# Function to process the image and generate caption and design
+@spaces.GPU
+def generate_caption_and_design(image):
+    # Generate caption
+    inputs = processor(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 design based on caption
+    generated_image = pipe(caption).images[0]
+    
+    return caption, generated_image
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=generate_caption_and_design,
+    inputs=gr.Image(type="pil", label="Upload an Image"),
+    outputs=[gr.Textbox(label="Generated Caption"), gr.Image(label="Generated Design")],
+    title="Image Caption and Design Generator",
+    description="Upload an image or provide an image URL to generate a caption and use it to create a similar design.",
+)
+
+# Launch Gradio app
+interface.launch()
 
-# Image upload or URL input
-img_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
-img_url = st.text_input("Or provide an image URL:")
-
-# Process the image
-raw_image = None
-if img_file:
-    raw_image = Image.open(img_file).convert("RGB")
-    st.image(raw_image, caption="Uploaded Image", use_column_width=True)
-elif img_url:
-    try:
-        raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
-        st.image(raw_image, caption="Image from URL", use_column_width=True)
-    except Exception as e:
-        st.error(f"Error loading image from URL: {e}")
-
-# Generate caption and design
-if raw_image and st.button("Generate Caption and Design"):
-    with st.spinner("Generating caption..."):
-        # 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)
-        st.success("Generated Caption:")
-        st.write(caption)
-
-    with st.spinner("Generating similar design..."):
-        # Generate similar design using the caption as a prompt
-        generated_image = pipe(caption).images[0]
-        st.success("Generated Design:")
-        st.image(generated_image, caption="Design Generated from Caption", use_column_width=True)