import gradio as gr
from gradio_client import Client, handle_file

from gradio_imageslider import ImageSlider

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(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 main(prompt, upscale_factor):
    step_one_flux = get_flux_image(prompt)
    step_two_upscale = get_upscale(prompt, step_one_flux, upscale_factor)
    return (step_one_flux, step_two_upscale)

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 Finegrained Image-Enhancer;")
        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
                )
                submit_btn = gr.Button("Submit")
        output_res = ImageSlider(label="Flux / Upscaled")

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

    submit_btn.click(
        fn=main,
        inputs=[prompt_in, upscale_factor],
        outputs=[output_res],
        
    )

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