from flask import Flask, request, jsonify
from gradio_client import Client

app = Flask(__name__)

# Initialize the Gradio Client
client = Client("stabilityai/stable-diffusion-3.5-large")

@app.route('/generate', methods=['POST'])
def generate_image():
    data = request.json

    # Extract parameters from the request
    prompt = data.get("prompt", "Hello!!")
    negative_prompt = data.get("negative_prompt", "")
    seed = data.get("seed", 0)
    randomize_seed = data.get("randomize_seed", True)
    width = data.get("width", 1024)
    height = data.get("height", 1024)
    guidance_scale = data.get("guidance_scale", 4.5)
    num_inference_steps = data.get("num_inference_steps", 40)

    try:
        # Use the Gradio Client to generate an image
        result = client.predict(
            prompt=prompt,
            negative_prompt=negative_prompt,
            seed=seed,
            randomize_seed=randomize_seed,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            api_name="/infer"
        )
        return jsonify({"image_url": result}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)