app.py

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

app = Flask(__name__)

# Load Stable Diffusion model from Hugging Face
model_id = "CompVis/stable-diffusion-v-1-4-original"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

@app.route('/generate-image', methods=['POST'])
def generate_image():
    description = request.json.get('description')
    if not description:
        return jsonify({'error': 'No description provided'}), 400
    
    # Generate image from description using Stable Diffusion
    image = pipe(description).images[0]
    
    # Save or serve the image
    image_path = save_image(image)  # Implement a function to save or serve the image
    
    return jsonify({'image_url': image_path})

def save_image(image):
    # Save the generated image to a path
    image_path = '/path/to/save/generated_image.png'
    image.save(image_path)
    return image_path

if __name__ == '__main__':
    app.run(debug=True)