import gradio as gr
import logging
import random
import warnings
import os
import numpy as np
import spaces
import torch
from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline
from gradio_imageslider import ImageSlider
from PIL import Image
from huggingface_hub import snapshot_download
from huggingface_hub import login
import subprocess

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

if torch.cuda.is_available():
    power_device = "GPU"
    device = "cuda"
else:
    power_device = "CPU"
    device = "cpu"

huggingface_token = os.getenv("token1704")

def login_to_huggingface():
    if huggingface_token:
        login(token=huggingface_token)
        model_path = snapshot_download(
            repo_id="black-forest-labs/FLUX.1-dev", 
            repo_type="model", 
            ignore_patterns=["*.md", "*..gitattributes"],
            local_dir="FLUX.1-dev",
            token=huggingface_token, # type a new token-id.
            )
        controlnet = FluxControlNetModel.from_pretrained(
                    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
        ).to(device)
        pipe = FluxControlNetPipeline.from_pretrained(
            model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
        )
        pipe.to(device)

MAX_SEED = 1000000
MAX_PIXEL_BUDGET = 1024 * 1024

def process_input(input_image, upscale_factor, **kwargs):
    w, h = input_image.size
    w_original, h_original = w, h
    aspect_ratio = w / h

    was_resized = False

    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
        warnings.warn(
            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels."
        )
        gr.Info(
            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing input to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels budget."
        )
        input_image = input_image.resize(
            (
                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
            )
        )
        was_resized = True

    # resize to multiple of 8
    w, h = input_image.size
    w = w - w % 8
    h = h - h % 8

    return input_image.resize((w, h)), w_original, h_original, was_resized


@spaces.GPU#(duration=42)
def infer(
    seed,
    randomize_seed,
    input_image,
    num_inference_steps,
    upscale_factor,
    controlnet_conditioning_scale,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    true_input_image = input_image
    input_image, w_original, h_original, was_resized = process_input(
        input_image, upscale_factor
    )

    # rescale with upscale factor
    w, h = input_image.size
    control_image = input_image.resize((w * upscale_factor, h * upscale_factor))

    generator = torch.Generator().manual_seed(seed)

    gr.Info("Upscaling image...")
    image = pipe(
        prompt="",
        control_image=control_image,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        num_inference_steps=num_inference_steps,
        guidance_scale=3.5,
        height=control_image.size[1],
        width=control_image.size[0],
        generator=generator,
    ).images[0]

    if was_resized:
        gr.Info(
            f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
        )

    # resize to target desired size
    image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
    image.save("output.jpg")
    
    # convert to numpy
    return [true_input_image, image, seed]

with gr.Blocks() as demo:

    with gr.Row():
        button = gr.LoginButton("Sign in")
        gr.load("models/black-forest-labs/FLUX.1-dev", accept_token=button, provider="hf-inference")
    
    # with gr.Row():
    #     run_button = gr.Button(value="Run")

    # with gr.Row():
    #     with gr.Column(scale=4):
    #         input_im = gr.Image(label="Input Image", type="pil")
    #     with gr.Column(scale=1):
    #         num_inference_steps = gr.Slider(
    #             label="Number of Inference Steps",
    #             minimum=8,
    #             maximum=50,
    #             step=1,
    #             value=28,
    #         )
    #         upscale_factor = gr.Slider(
    #             label="Upscale Factor",
    #             minimum=1,
    #             maximum=4,
    #             step=1,
    #             value=4,
    #         )
    #         controlnet_conditioning_scale = gr.Slider(
    #             label="Controlnet Conditioning Scale",
    #             minimum=0.1,
    #             maximum=1.5,
    #             step=0.1,
    #             value=0.6,
    #         )
    #         seed = gr.Slider(
    #             label="Seed",
    #             minimum=0,
    #             maximum=MAX_SEED,
    #             step=1,
    #             value=42,
    #         )

    #         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    # with gr.Row():
    #     result = ImageSlider(label="Input / Output", type="pil", interactive=True)

    # examples = gr.Examples(
    #     examples=[
    #         [42, False, "z1.webp", 28, 4, 0.6],
    #         [42, False, "z2.webp", 28, 4, 0.6],
 
    #     ],
    #     inputs=[
    #         seed,
    #         randomize_seed,
    #         input_im,
    #         num_inference_steps,
    #         upscale_factor,
    #         controlnet_conditioning_scale,
    #     ],
    #     fn=infer,
    #     outputs=result,
    #     cache_examples="lazy",
    # )

    # gr.on(
    #     [run_button.click],
    #     fn=infer,
    #     inputs=[
    #         seed,
    #         randomize_seed,
    #         input_im,
    #         num_inference_steps,
    #         upscale_factor,
    #         controlnet_conditioning_scale,
    #     ],
    #     outputs=result,
    #     show_api=False,
    #     # show_progress="minimal",
    # )
    login_to_huggingface()

demo.queue().launch(share=True, show_api=True)

@app.get("/")
def read_root():
    return {"message": "Hello, World!"}

@app.get("/items/{item_id}")
def read_item(item_id: int, q: str = None):
    return {"item_id": item_id, "q": q}