app.py

from diffusers import StableDiffusionPipeline
import torch
from flask import Flask, request, jsonify

app = Flask(__name__)

# Load the model
model_id = "ZB-Tech/Text-to-Image"
pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipeline.to("cuda" if torch.cuda.is_available() else "cpu")

@app.route("/generate", methods=["POST"])
def generate_image():
    data = request.get_json()
    prompt = data.get("prompt", "A scenic landscape")
    
    image = pipeline(prompt).images[0]
    image_path = "generated_image.png"
    image.save(image_path)
    
    return jsonify({"image_url": image_path})

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