from diffusers import EulerAncestralDiscreteScheduler, StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
from PIL import Image
import gradio as gr
import random
import torch
import math
import spaces

device = "cuda" if torch.cuda.is_available() else "cpu"
negative_prompt = "deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck"

css = """
.btn-green {
  background-image: linear-gradient(to bottom right, #6dd178, #00a613) !important;
  border-color: #22c55e !important;
  color: #166534 !important;
}
.btn-green:hover {
  background-image: linear-gradient(to bottom right, #6dd178, #6dd178) !important;
}
"""

@spaces.GPU(duration=60, enable_queue=True)
def generate(prompt, samp_steps, seed, progress=gr.Progress(track_tqdm=True)):
    print("prompt = ", prompt)
    turbo_steps=5 # use sweet spot
    if seed < 0:
        seed = random.randint(1,999999)
    image = txt2img(
        prompt=prompt,
		negative_prompt=negative_prompt,
        num_inference_steps=turbo_steps,  
        guidance_scale=2,
        seed=seed,
    ).images[0]
    upscaled_image = image.resize((1024,1024), 1)
    final_image = img2img(
        prompt=prompt,
		negative_prompt=negative_prompt,
        image=upscaled_image,
        num_inference_steps=max(1,samp_steps-turbo_steps*2), # always discount lightning at 2x efficiency
        guidance_scale=5,
        strength=1,
        seed=seed,
    ).images[0]
    return [final_image], seed


def set_base_models():
    txt2img = StableDiffusionXLPipeline.from_pretrained(
        "SG161222/RealVisXL_V5.0_Lightning",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        use_safetensors=True,
        add_watermarker=False,
    ).to(device)
    txt2img.scheduler = EulerAncestralDiscreteScheduler.from_config(txt2img.scheduler.config)
	
    img2img = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        "SG161222/RealVisXL_V5.0",
        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32,
		use_safetensors=True,
        add_watermarker=False,
    ).to(device)
    img2img.scheduler = EulerAncestralDiscreteScheduler.from_config(img2img.scheduler.config)

    return txt2img, img2img

with gr.Blocks(css=css) as demo:
    with gr.Column():
        prompt = gr.Textbox(label="Prompt for RealVisXL_V5")
        submit_btn = gr.Button("Generate", elem_classes="btn-green")
        
        with gr.Row():
            turbo_steps = gr.Slider(5, 5, value=5, step=1, label="Turbo steps (fixed at 5)")
            sampling_steps = gr.Slider(15, 50, value=25, step=1, label="Effective steps")
            seed = gr.Number(label="Seed", value=-1, minimum=-1, precision=0)
            lastSeed = gr.Number(label="Last Seed", value=-1, interactive=False)
            
        gallery = gr.Gallery(show_label=False, preview=True, container=False, height=1100)
        
    submit_btn.click(generate, [prompt, sampling_steps, seed], [gallery, lastSeed], queue=True)
    
txt2img, img2img = set_base_models()
demo.launch(debug=True, share=True)