Update app.py

app.py

@@ -1,54 +1,51 @@
 import os
-import requests
-from PIL import Image
-import torch
-import gradio as gr
 from huggingface_hub import login
-from transformers import AutoProcessor, AutoModelForCausalLM
-from diffusers import DiffusionPipeline
+from transformers import BlipProcessor, BlipForConditionalGeneration
+
+
 
-# Hugging Face token setup
+# Get Hugging Face Token from environment variable
 hf_token = os.getenv('HF_AUTH_TOKEN')
 if not hf_token:
     raise ValueError("Hugging Face token is not set in the environment variables.")
 login(token=hf_token)
 
-# Initialize Stable Diffusion pipeline
-pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")
+# Load the processor and model
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+import gradio as gr
+from diffusers import DiffusionPipeline
+import torch
+import spaces  # Hugging Face Spaces module
 
-# Initialize captioning model and processor
-caption_model_name = "pretrained-caption-model"  # Replace with the actual model name
-processor = AutoProcessor.from_pretrained(caption_model_name)
-model = AutoModelForCausalLM.from_pretrained(caption_model_name)
-
-# 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)
-
-# 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()
+# 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)