# import random
# from flask import Flask, request, jsonify, send_file
# from flask_cors import CORS
from diffusers import StableDiffusionPipeline, AutoPipelineForText2Image
# import torch
from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
import requests
import accelerate

# app = Flask(__name__)
app = FastAPI()
# CORS(app)

origins = ["https://felix-ng.com"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_methods=["*"],
    allow_headers=["*"],
)

pipeline = AutoPipelineForText2Image.from_pretrained(
"stabilityai/stable-diffusion-2", variant="fp16", use_safetensors=True
  )

@app.get('/')
def read_root():
    return {"Hello": "World!"}

@app.get('/image/{prompt}')
def generate_image(prompt):
    print(prompt)
  
    image = pipeline(prompt=prompt, num_inference_steps=3).images[0]
    print("Saving image")
    image.save("image_generated.png")

  
  # buffered = BytesIO()
  # image.save(buffered, format="PNG")
  # encoded_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
  # return jsonify({'image_generated_path': image_generated_path}), 201
  # return jsonify({'image_generated_path': image_generated_path, 'image_base64': encoded_image}), 201
    return FileResponse('image_generated.png')


# if __name__ == '__main__':
#   host = '127.0.0.1'
#   port = 5000
#   app.run(debug=True)