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)