Upload 2 files

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+from PIL import Image
+from diffusers import StableDiffusionImg2ImgPipeline
+import torch
+import uuid
+
+# Set device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+# Load model
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=dtype,
+    use_safetensors=True
+).to(device)
+
+# Resize to match selected aspect ratio
+def resize_to_aspect(image, aspect_ratio):
+    width, height = image.size
+    aspect_map = {
+        "1:1": (min(width, height), min(width, height)),
+        "16:9": (width, int(width * 9 / 16)),
+        "4:5": (width, int(width * 5 / 4)),
+        "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))
+
+def resize_to_512(image):
+    return image.resize((512, 512))
+
+# Generate the new 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)
+    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
+
+# Launch interface
+demo = gr.Interface(
+    fn=generate_img,
+    inputs=[
+        gr.Image(type="pil", label="Upload Product Image", image_mode='RGB'),
+        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')
+    ],
+    title="Image-to-Image Product Generator",
+    description="Upload a product image, describe your idea, and select the output aspect ratio."
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True, share=True, ssr_mode=False)

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+transformers
+diffusers
+safetensors