Update app.py

app.py

@@ -2,21 +2,19 @@ import gradio as gr
 from PIL import Image
 from diffusers import StableDiffusionImg2ImgPipeline
 import torch
-import uuid
 
-# Set device
+# Detect device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-# dtype = torch.float16 if device == "cuda" else torch.float32
-dtype = torch.float32
+dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Load model
+# Load the model with appropriate dtype
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     torch_dtype=dtype,
     use_safetensors=True
 ).to(device)
 
-# Resize to match selected aspect ratio
+# Resize uploaded image to selected aspect ratio
 def resize_to_aspect(image, aspect_ratio):
     width, height = image.size
     aspect_map = {
@@ -26,36 +24,32 @@ def resize_to_aspect(image, aspect_ratio):
         "9:16": (int(height * 9 / 16), height)
     }
     target_w, target_h = aspect_map.get(aspect_ratio, (width, height))
-    return image.resize((target_w, target_h))
+    image = image.resize((target_w, target_h))
+    return image
 
 def resize_to_512(image):
     return image.resize((512, 512))
 
-# Generate the new image
+
+#
+# Generate new image from prompt and reference image
 def generate_img(product_img, prompt, aspect_ratio):
     resized_img = resize_to_aspect(product_img, aspect_ratio).convert("RGB")
-    resized_img = resize_to_512(resized_img)
+    resized_img = resize_to_512(resized_img)  # Required by the model
     output = pipe(prompt=prompt, image=resized_img, strength=0.75, guidance_scale=7.5)
-    image = output.images[0]
-    save_path = f"/tmp/generated_{uuid.uuid4().hex}.png"
-    image.save(save_path)
-    return image, save_path
+    return output.images[0]
 
-# Launch interface
-app = gr.Interface(
+# Gradio UI
+gr.Interface(
     fn=generate_img,
     inputs=[
-        gr.Image(type="pil", label="Upload Product Image", image_mode='RGB'),
+        gr.Image(type="pil", label="Upload Product Image"),
         gr.Textbox(label="Prompt", placeholder="Describe what you want to generate"),
         gr.Dropdown(["1:1", "16:9", "4:5", "9:16"], label="Aspect Ratio", value="1:1")
     ],
-    outputs=[
-        gr.Image(label='Preview'),
-        gr.File(label='Download Image')
-    ],
+    # outputs=gr.Image(label="Generated Image", type="pil"),
+    outputs=gr.Image(label='Download Image'),
     title="Image-to-Image Product Generator",
-    description="Upload a product image, describe your idea, and select the output aspect ratio."
-)
+    description="upload a product image, describe your idea, and select the output aspect ratio."
 
-# if __name__ == "__main__":
-#     app.launch(server_name="0.0.0.0", server_port=7860, debug=True, ssr_mode=False)
+).launch()