from diffusers import DiffusionPipeline
import gradio as gr
import numpy as np
import random
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16

pipe = DiffusionPipeline.from_pretrained("Chan-Y/Stable-Flash-Lightning", 
                                             torch_dtype=torch.float16).to(device)

MAX_SEED = 999999999999999
MAX_IMAGE_SIZE = 1344

def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    generator = torch.Generator().manual_seed(seed)
    
    image = pipe(
        prompt = prompt, 
        negative_prompt = negative_prompt,
        guidance_scale = guidance_scale, 
        num_inference_steps = num_inference_steps, 
        width = width, 
        height = height,
        generator = generator
    ).images[0] 
    
    return image, seed


examples = [
    ["high quality, 8K Ultra HD, high detailed, masterpiece, A digital illustration of anime style, soft anime tones, Detailed illustration of many colorful soap bubbles falling from the sky on a beautiful woman, colorful colors, colorful woman, pink hair, blue eye, luminism, three dimensional effect, enhanced beauty, Albert Anker, Kyoto Animation, Greg Rutkowski, Artgerm, WLOP, Alphonse Beeple, luminism, 3d render, octane render, Isometric, by yukisakura, awesome full color,",
     "nude, nsfw, text, letters, too many feet, too many fingers, long neck, 2 heads, duplicate, abstract, disfigured, deformed, toy, figure, framed, disfigured, bad art,deformed, poorly drawn, extra limbs, weird colors, 2 heads, elongated body, cropped image, out of frame, draft, deformed hands, twisted fingers, double image, malformed hands, multiple heads, extra limb, ugly, poorly drawn hands, missing limb, cut-off, over satured, grain, lowères, bad anatomy, poorly drawn face, mutation, mutated, floating limbs, disconnected limbs, out of focus, long body, disgusting, extra fingers, groos proportions, missing arms, mutated hands, cloned face, missing legs,", 
     728456274845100,
     28],
    ["cute catgirl , orange and white, with furry ears and striped tail, devious expression, full-body shot, 8K.", 
     "Double limbs, cross eyes, long limbs, deformed fingers, fat, ugly, missing body parts, pixelation", 
     712099815579800,
     50],
    ["pink female witch portrait beauty pretty soft dewy glow skin sparkle highlighter pink haired youthful witch web witch feminine retro vintage antique spiderwebs on her pointed pink witchhat wearing a spiderweb pink corset ribbon laced up, with dewdrops on a pink roses with spiderwebs with dewdrops on the webs halloween spiderwebs spider-web hd highres 4k 8k coquette aesthetic ",
     "",       
     711753355540400,
     50]
]


css="""
#col-container {
    margin: 0 auto;
    max-width: 580px;
}
"""


with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # Demo [Chan-Y/Stable-Flash-Lightning](https://huggingface.co/Chan-Y/Stable-Flash-Lightning)
        by Cihan Yalçın | My [LinkedIn](https://www.linkedin.com/in/chanyalcin/) My [GitHub](https://github.com/g-hano)
        """)
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
            )
            
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )
            
            with gr.Row():
                
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=5.0,
                )
                
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
    with gr.Accordion("Examples", open=False):
        gr.Examples(
            examples=examples,
            inputs=[prompt, negative_prompt, seed, num_inference_steps]
        )
    gr.on(
        triggers=[run_button.click, prompt.submit, negative_prompt.submit],
        fn = infer,
        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs = [result, seed]
    )

demo.launch()