Update app.py

app.py

@@ -1,20 +1,38 @@
 import torch
 from diffusers import AutoPipelineForImage2Image
-from diffusers.utils import make_image_grid, load_image
+from PIL import Image
+import gradio as gr
 
+# Initialize the image-to-image pipeline
 pipeline = AutoPipelineForImage2Image.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True
 )
+
+# Enable CPU offloading to manage memory usage
 pipeline.enable_model_cpu_offload()
-# remove following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
-# pipeline.enable_xformers_memory_efficient_attention()
 
-# prepare image
-url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-init.png"
-init_image = load_image(url)
+# Function to process the image and prompt, and generate a new image
+def process_image(prompt, init_image):
+    # Convert the uploaded file to an image
+    init_image = Image.open(init_image).convert("RGB")
+
+    # Generate an image based on the prompt and initial image
+    with torch.no_grad():  # This ensures that no gradients are calculated, saving memory
+        generated_image = pipeline(prompt, init_image=init_image, num_inference_steps=50).images[0]
 
-prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
+    return generated_image
+
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=process_image,
+    inputs=[gr.inputs.Textbox(label="Prompt"), gr.inputs.Image(label="Initial Image", type="pil")],
+    outputs=gr.outputs.Image(type="pil", label="Generated Image"),
+    title="Image-to-Image Transformation",
+    description="Upload an image and enter a prompt to transform the image accordingly."
+)
 
-# pass prompt and image to pipeline
-image = pipeline(prompt, image=init_image).images[0]
-make_image_grid([init_image, image], rows=1, cols=2)
+# Launch the interface
+interface.launch()