import os
import time
import sys
import re
import random
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()}")

		if '_seed' in data.string:
			seed = 1024
		else:
		 	seed = random.randrange(1024)

		text = re.sub('^#', '', data.string)
		text = re.sub('_seed', '', text)

		# prompt = '(('+text+')) (( 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 = '(('+text+')) (( 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')

		generator = torch.Generator("cuda").manual_seed(seed)
		image = pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, generator=generator).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": ""}

app.mount("/static", StaticFiles(directory="/code/tmpdir"), name="/static")