Update app.py

app.py

@@ -1,48 +1,26 @@
-from flask import Flask, request, jsonify
-from diffusers import DiffusionPipeline
+import gradio as gr
 import torch
-from PIL import Image
-import io
-import base64
-
-app = Flask(__name__)
-
-# Load the instruct-pix2pix model
-model_id = "timbrooks/instruct-pix2pix"
-pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipeline.to("cuda")  # Use "cpu" if you're running without a GPU
-
-@app.route("/", methods=["GET"])
-def home():
-    return "Welcome to the Instruct-Pix2Pix API!"
-
-@app.route("/edit-image", methods=["POST"])
-def edit_image():
-    try:
-        # Extract the prompt and image from the request
-        data = request.json
-        prompt = data.get("prompt", "A beautiful landscape with a sunset")
-        image_data = data.get("image")  # Expected as base64 encoded string
-
-        # Decode base64 image
-        image = Image.open(io.BytesIO(base64.b64decode(image_data)))
-
-        # Run the model with the prompt and image
-        edited_image = pipeline(prompt=prompt, image=image).images[0]
-
-        # Save the edited image
-        output_image_path = "edited_image.png"
-        edited_image.save(output_image_path)
-
-        # Optionally return the image as base64 in the response
-        buffered = io.BytesIO()
-        edited_image.save(buffered, format="PNG")
-        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-        return jsonify({"message": "Image edited successfully!", "edited_image": img_str})
-
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+from diffusers import DiffusionPipeline
 
+# Load the pre-trained model from Hugging Face
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16)
+pipe.to("cuda")  # Ensure the model runs on the GPU if available
+
+# Define the function for the Gradio interface
+def generate_image(prompt):
+    # Generate an image using the provided prompt
+    image = pipe(prompt).images[0]
+    return image
+
+# Set up the Gradio interface
+interface = gr.Interface(
+    fn=generate_image,
+    inputs="text",
+    outputs="image",
+    title="Stable Diffusion XL Refiner",
+    description="Generate images from text prompts using Stable Diffusion XL Refiner 1.0"
+)
+
+# Launch the Gradio app
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000)
+    interface.launch()