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

app = Flask(__name__)

# Load the model
model_id = "kothariyashhh/GenAi-Texttoimage"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Use float32 if running on CPU
torch_dtype = torch.float16 if device == "cuda" else torch.float32

pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
pipeline.to(device)

@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)