import os
import time
import sys
import torch
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from pydantic import Field
from diffusers import StableDiffusionPipeline

app = FastAPI()

class Data(BaseModel):
	string: str
	member_secret: str

class ItemOut(BaseModel):
	status: str
	file: str

@app.get("/")
def index():
	return "SORRY! This file is member only."

@app.post("/draw", response_model=ItemOut)
def draw(data: Data):
	if data.member_secret != "" and data.member_secret == os.environ.get("MEMBER_SECRET"):
		print(f"Is CUDA available: {torch.cuda.is_available()}")

		# prompt = '(('+data.string+')) (( photograph )), highly detailed, sharp focus, 8k, 4k, (( photorealism )), detailed, saturated, portrait, 50mm, F/2.8, 1m away, ( global illumination, studio light, volumetric light ), ((( multicolor lights )))'
		prompt = '(('+data.string+')) (( photograph )), highly detailed, sharp focus, 8k, 4k, (( photorealism )), detailed, saturated, portrait, 50mm, F/2.8, 1m away, ((( multicolor lights )))'
		n_prompt = 'text, blurry, art, painting, rendering, drawing, sketch, (( ugly )), (( duplicate )), ( morbid ), (( mutilated )), ( mutated ), ( deformed ), ( disfigured ), ( extra limbs ), ( malformed limbs ), ( missing arms ), ( missing legs ), ( extra arms ), ( extra legs ), ( fused fingers ), ( too many fingers ), long neck, low quality, worst quality'

		# https://huggingface.co/docs/hub/spaces-sdks-docker-first-demo
		# how to validation: https://qiita.com/bee2/items/75d9c0d7ba20e7a4a0e9
		# https://github.com/huggingface/diffusers

		model_id = 'stabilityai/stable-diffusion-2'

		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
		pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
		pipe = pipe.to('cuda')

		image = pipe(prompt, negative_prompt=n_prompt).images[0]

		fileName = "sd_" + str(time.time()) + '.png'
		image.save("/code/tmpdir/" + fileName)

		print(fileName)

		return {"status": "OK", "file": fileName}
	else:
		return {"status": "SORRY! This file is member only.", "file": ""}