app.py

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)