Create app.py

app.py

@@ -0,0 +1,48 @@
+from flask import Flask, request, jsonify
+from diffusers import DiffusionPipeline
+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
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000)