|
from flask import Flask, request |
|
from diffusers import StableDiffusionPipeline |
|
import torch |
|
import os |
|
|
|
app = Flask(__name__) |
|
|
|
|
|
model_id = "runwayml/stable-diffusion-v1-5" |
|
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) |
|
pipe = pipe.to("cpu") |
|
|
|
def dummy(images, **kwargs): |
|
return images, False |
|
|
|
pipe.safety_checker = dummy |
|
|
|
@app.route('/') |
|
def generate_image(): |
|
prompt = request.args.get('prompt') |
|
image = pipe(prompt).images[0] |
|
|
|
image_data = image.tobytes().hex() |
|
|
|
return {'image_data': image_data} |
|
|
|
if __name__ == '__main__': |
|
app.run(host='0.0.0.0', port=5000) |
