from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

import subprocess
import os
import json
import uuid

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler

app = FastAPI() 

@app.get("/generate")
def generate_image(prompt, inference_steps, model):
    torch.cuda.empty_cache()

    pipeline = StableDiffusionPipeline.from_pretrained(str(model), torch_dtype=torch.float16)
    pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
    pipeline = pipeline.to("cuda")
 
    image = pipeline(prompt, num_inference_steps=int(inference_steps), height=512, width=512).images[0]

    filename = str(uuid.uuid4()) + ".jpg"
    image.save(filename)
    
    assertion = {
               "assertions": [
                   {
                   "label": "com.truepic.custom.ai",
                   "data": {
                       "model_name": model,
                       "model_version": "1.0",
                       "prompt": prompt
                   }
                   }
               ]
               }

    json_object = json.dumps(assertion, indent=4)
    with open("assertion.json", "w") as outfile:
       outfile.write(json_object)

    subprocess.check_output(['./truepic-sign', 'init', 'file-system', '--api-key', os.environ.get("api_key")])
    subprocess.check_output(['./truepic-sign', 'sign', filename, '--assertions-file', 'assertion.json', '--output', (os.getcwd() + '/static/' + filename)])
    
    return {"response": filename}


app.mount("/", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/static/index.html", media_type="text/html")