import gradio as gr
import os 
from gradio_client import Client, handle_file
from huggingface_hub import login

from gradio_imageslider import ImageSlider

hf_tkn = os.environ.get("HF_TKN")
login(hf_tkn)

def get_flux_image(prompt):
    client = Client("black-forest-labs/FLUX.1-schnell")
    result = client.predict(
		prompt=prompt,
		seed=0,
		randomize_seed=True,
		width=1024,
		height=1024,
		num_inference_steps=4,
		api_name="/infer"
    )
    print(result)
    return result[0]

def get_upscale_finegrain(prompt, img_path, upscale_factor):
    client = Client("finegrain/finegrain-image-enhancer")
    result = client.predict(
		input_image=handle_file(img_path),
		prompt=prompt,
		negative_prompt="",
		seed=42,
		upscale_factor=upscale_factor,
	    controlnet_scale=0.6,
		controlnet_decay=1,
		condition_scale=6,
		tile_width=112,
		tile_height=144,
		denoise_strength=0.35,
		num_inference_steps=18,
		solver="DDIM",
		api_name="/process"
    )
    print(result)
    return result[1]

def get_clarity_upscale(prompt, img_path, upscale_factor):
    client = Client("jbilcke-hf/clarity-upscaler")
    result = client.predict(
		img_path,	# filepath  in 'Image' Image component
		prompt,	# str  in 'Prompt' Textbox component
		"",	# str  in 'Negative Prompt' Textbox component
		upscale_factor,	# float  in 'Scale Factor' Number component
		1,	# float (numeric value between 1 and 50) in 'Dynamic' Slider component
		3,	# float  in 'Creativity' Number component
		3,	# float  in 'Resemblance' Number component
		"16",	# Literal['16', '32', '48', '64', '80', '96', '112', '128', '144', '160', '176', '192', '208', '224', '240', '256']  in 'tiling_width' Dropdown component
		"16",	# Literal['16', '32', '48', '64', '80', '96', '112', '128', '144', '160', '176', '192', '208', '224', '240', '256']  in 'tiling_height' Dropdown component
		"epicrealism_naturalSinRC1VAE.safetensors [84d76a0328]",	# Literal['epicrealism_naturalSinRC1VAE.safetensors [84d76a0328]', 'juggernaut_reborn.safetensors [338b85bc4f]', 'flat2DAnimerge_v45Sharp.safetensors']  in 'sd_model' Dropdown component
		"DPM++ 2M Karras",	# Literal['DPM++ 2M Karras', 'DPM++ SDE Karras', 'DPM++ 2M SDE Exponential', 'DPM++ 2M SDE Karras', 'Euler a', 'Euler', 'LMS', 'Heun', 'DPM2', 'DPM2 a', 'DPM++ 2S a', 'DPM++ 2M', 'DPM++ SDE', 'DPM++ 2M SDE', 'DPM++ 2M SDE Heun', 'DPM++ 2M SDE Heun Karras', 'DPM++ 2M SDE Heun Exponential', 'DPM++ 3M SDE', 'DPM++ 3M SDE Karras', 'DPM++ 3M SDE Exponential', 'DPM fast', 'DPM adaptive', 'LMS Karras', 'DPM2 Karras', 'DPM2 a Karras', 'DPM++ 2S a Karras', 'Restart', 'DDIM', 'PLMS', 'UniPC']  in 'scheduler' Dropdown component
		1,	# float (numeric value between 1 and 100) in 'Num Inference Steps' Slider component
		3,	# float  in 'Seed' Number component
		True,	# bool  in 'Downscaling' Checkbox component
		3,	# float  in 'Downscaling Resolution' Number component
		"Hello!!",	# str  in 'Lora Links' Textbox component
		"Hello!!",	# str  in 'Custom Sd Model' Textbox component
		api_name="/predict"
    )
    print(result)
    return result

def main(prompt, upscale_factor, upscale_provider):
    step_one_flux = get_flux_image(prompt)
    if upscale_provider == "finegrain image enhancer":
        step_two_upscale = get_upscale_finegrain(prompt, step_one_flux, upscale_factor)
    elif upscale_provider == "clarity upscale":
        step_two_upscale = get_clarity_upscale(prompt, step_one_flux, upscale_factor)
    return (step_one_flux, step_two_upscale)

def clean_previous():
    return gr.update(value=None)

css = """
#col-container{
    margin: 0 auto;
    max-width: 1024px;
}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# Flux Upscaled")
        gr.Markdown("Step 1: Generate image with FLUX schnell; Step 2: UpScale with Finegrain Image-Enhancer OR Clarity UpScale;")
        with gr.Group():
            prompt_in = gr.Textbox(label="Prompt")
            with gr.Row():
                upscale_factor = gr.Radio(
                    label = "UpScale Factor",
                    choices = [
                        2, 3, 4
                    ],
                    value = 2,
                    scale=2
                )
                upscale_provider = gr.Dropdown(
                    label = "UpScale Provider",
                    choices = ["finegrain image enhancer", "clarity upscale"],
                    value = "clarity upscale",
                    scale=2
                )
                submit_btn = gr.Button("Submit", scale=1)
        output_res = ImageSlider(label="Flux / Upscaled")

        gr.Examples(
            examples = [
                ["a tiny astronaut hatching from an egg on the moon", 2, "clarity upscale"],
                ["a bright blue bird in the garden, natural photo cinematic, MM full HD", 2, "clarity upscale"]
            ],
            fn = main,
            inputs=[prompt_in, upscale_factor, upscale_provider],
            outputs=[output_res],
            cache_examples = "lazy"
        )

    submit_btn.click(
        fn = clean_previous,
        inputs = None,
        outputs = [output_res],
        queue=False
    ).then(
        fn=main,
        inputs=[prompt_in, upscale_factor, upscale_provider],
        outputs=[output_res],
        
    )

demo.queue().launch(show_api=False, show_error=True)