from config import *
from diffusers import AutoPipelineForText2Image
import torch
from argparse import ArgumentParser
import humanize
import datetime as dt

def generate_image(path, imgfile, prompt):
	model_id = "stabilityai/sdxl-turbo"
	pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16)
	pipe = pipe.to("cuda")
	prompt = prompt
	image = pipe(prompt,num_inference_steps=1, guidance_scale=0.0).images[0]
	image.save("images\output.png")
	
	# pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
	# image = pipe(prompt=prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
	# image.save(os.path.join(path, imgfile))
	#image.save(path)


if __name__ == '__main__':
	path_id = str(int(time.time()))
	path = os.path.join("temp", "image", path_id)
	os.makedirs(path, exist_ok=True)

	parser = ArgumentParser()
	parser.add_argument("--prompt", default="Young asian. professional woman with long, blonde hair, smiling slightly", 
					 help="avatar prompt")
	args = parser.parse_args()

	tstart = time.time()	

	generate_image(path, "avatar.png", 
				f"hyper-realistic photo, centered, {args.prompt}, \
				rim lighting, studio lighting, looking at the camera, front ")

	print("total time:", humanize.naturaldelta(dt.timedelta(seconds=int(time.time() - tstart))))