# credits : https://huggingface.co/spaces/black-forest-labs/FLUX.1-dev

import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images

hf_token = os.getenv("HF_TOKEN")

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

taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
torch.cuda.empty_cache()

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

@spaces.GPU(duration=75)

def infer(name, pet, background, style, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    if pet == "Kaatz":
        intro = "please generate an image of a cat sitting "
    elif pet == "Mupp":
        intro = "please generate an image of a dog sitting "
    elif pet == "Hues":
        intro = "please generate an image of a bunny sitting "
    else:
        intro = "please generate an image of an hamster sitting "

    if background == "Wunnzëmmer":
        place = intro + "in a living space "
    elif background == "Grafitti Mauer":  
        place = intro + "in front of a wall with graffiti "
    elif background == "Strooss":   
        place = intro + "in a street in the city "
    elif background == "Plage":
        place = intro + "at the beach "
    else:
        place = intro + " in the forest "

    if style == "Photo":
        prompt = place + "holding a signal that says " + name + "in a photorealistic style"
    elif style == "Cartoon":
        prompt = place + "holding a signal that says " + name + "in a cartoon style"
    elif style == "Woll":
        prompt = place + "holding a signal that says " + name + "in a knitted with wool style"
    elif style == "Aquarell":
        prompt = place + "holding a signal that says " + name + "in a watercolorl style"
    else:
        prompt = place + "holding a signal that says " + name + "in a 3D style"
        
    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt=prompt,
            guidance_scale=4,
            num_inference_steps=28,
            width=1024,
            height=1024,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,
        ):
            yield img, seed

examples = [
    "a tiny astronaut hatching from an egg on the moon",
    "a cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
]

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

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# Mäin éischt KI-Bild
        Mol mer e Bild mat mengem Hausdéier a mengem Numm op engem Schëld !
        """) 
        
        with gr.Row():
            
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Schreif däin Text mat dengem Numm",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)

        with gr.Row():
            pet = gr.Radio(
                choices=["Kaatz", "Mupp", "Hues", "Hamster"],
                label="Hausdéier",
                value="Kaatz"
            ) 

        with gr.Row():
            background = gr.Radio(
                choices=["Wunnzëmmer", "Grafitti Mauer", "Strooss", "Plage", "Bësch"],
                label="Hannergrond",
                value="Strooss"
            ) 

        with gr.Row():
            style = gr.Radio(
                choices=["Photo", "Cartoon", "Woll", "Aquarell", "3D"],
                label="Style",
                value="Photo"
            ) 
           
        result = gr.Image(label="Result", show_label=False) 

        with gr.Accordion("Advanced Settings", open=False):
            
            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=32,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples="lazy"
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, pet, background, style],
        outputs = [result, seed]
    )

demo.launch()