Spaces:
Runtime error
Runtime error
File size: 1,097 Bytes
ba11222 f977a8b 395adf9 ba11222 f977a8b 395adf9 f977a8b ba11222 f977a8b ba11222 395adf9 ba11222 bd7827c ba11222 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
import os
from flask import Flask, request, render_template, send_file
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt
# Find models in https://huggingface.co/models?pipeline_tag=text-to-image&library=diffusers&sort=trending
model_id = "stabilityai/stable-diffusion-2-1"
imagesPath = "images"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to("cpu")
app = Flask("AI API")
@app.get("/")
def read_root():
return render_template("index.html")
@app.route("/api", methods=["POST"])
def receive_data():
data = request.get_json()
print("Prompt:", data["prompt"])
prompt = data["prompt"]
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
image = pipe(prompt).images[0]
# Convert the torch Tensor to a NumPy array and move to CPU
image_np = image.cpu().numpy()
print("[Prompt]: ", prompt)
plt.imsave(f"{imagesPath}/{prompt}.png", image_np.transpose(1, 2, 0))
return send_file(os.path.join(imagesPath, f"{prompt}.png"), mimetype='image/png')
app.run(host="0.0.0.0", port=7860, debug=False) |