import sys
sys.path.append('./')

import gradio as gr
import spaces
import os
import sys
import subprocess
import numpy as np
from PIL import Image
import cv2
import torch
import random

os.system("pip install -e ./controlnet_aux")

from controlnet_aux import OpenposeDetector #, CannyDetector
from depth_anything_v2.dpt import DepthAnythingV2

from huggingface_hub import hf_hub_download

from huggingface_hub import login
hf_token = os.environ.get("HF_TOKEN")
login(token=hf_token)

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

try:
    local_dir = os.path.dirname(__file__)
except:
    local_dir = '.'
    
hf_hub_download(repo_id="briaai/BRIA-3.1", filename='pipeline_bria.py', local_dir=local_dir)
hf_hub_download(repo_id="briaai/BRIA-3.1", filename='transformer_bria.py', local_dir=local_dir)
hf_hub_download(repo_id="briaai/BRIA-3.1", filename='bria_utils.py', local_dir=local_dir)
hf_hub_download(repo_id="briaai/BRIA-3.1-ControlNet-Union", filename='pipeline_bria_controlnet.py', local_dir=local_dir)
hf_hub_download(repo_id="briaai/BRIA-3.1-ControlNet-Union", filename='controlnet_bria.py', local_dir=local_dir)

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model_configs = {
    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
}

RATIO_CONFIGS_1024 = {
    0.6666666666666666: {"width": 832, "height": 1248},
    0.7432432432432432: {"width": 880, "height": 1184},
    0.8028169014084507: {"width": 912, "height": 1136},
    1.0: {"width": 1024, "height": 1024},
    1.2456140350877194: {"width": 1136, "height": 912},
    1.3454545454545455: {"width": 1184, "height": 880},
    1.4339622641509433: {"width": 1216, "height": 848},
    1.5: {"width": 1248, "height": 832},
    1.5490196078431373: {"width": 1264, "height": 816},
    1.62: {"width": 1296, "height": 800},
    1.7708333333333333: {"width": 1360, "height": 768},
}

encoder = 'vitl'
model = DepthAnythingV2(**model_configs[encoder])
filepath = hf_hub_download(repo_id=f"depth-anything/Depth-Anything-V2-Large", filename=f"depth_anything_v2_vitl.pth", repo_type="model")
state_dict = torch.load(filepath, map_location="cpu")
model.load_state_dict(state_dict)
model = model.to(DEVICE).eval()

import torch
from diffusers.utils import load_image
from controlnet_bria import BriaControlNetModel, BriaMultiControlNetModel
from pipeline_bria_controlnet import BriaControlNetPipeline
import PIL.Image as Image

base_model = 'briaai/BRIA-3.1'
controlnet_model = 'briaai/BRIA-3.1-ControlNet-Union'
controlnet = BriaControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
pipe = BriaControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16, trust_remote_code=True)
pipe = pipe.to(device="cuda", dtype=torch.bfloat16)

mode_mapping = {
    "depth": 0,
    "canny": 1,
    "colorgrid": 2,
    "recolor": 3,
    "tile": 4,
    "pose": 5,
}
strength_mapping = {
    "depth": 1.0,
    "canny": 1.0,
    "colorgrid": 1.0,
    "recolor": 1.0,
    "tile": 1.0,
    "pose": 1.0,
}
open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")

torch.backends.cuda.matmul.allow_tf32 = True
pipe.enable_model_cpu_offload() # for saving memory

def convert_from_image_to_cv2(img: Image) -> np.ndarray:
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

def convert_from_cv2_to_image(img: np.ndarray) -> Image:
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

def extract_depth(image):
    image = np.asarray(image)
    depth = model.infer_image(image[:, :, ::-1])
    depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
    depth = depth.astype(np.uint8)
    gray_depth = Image.fromarray(depth).convert('RGB') 
    return gray_depth

def extract_openpose(img):
    processed_image_open_pose = open_pose(img, hand_and_face=True)
    processed_image_open_pose = processed_image_open_pose.resize(img.size)
    return processed_image_open_pose
    
def extract_canny(input_image):
    image = np.array(input_image)
    image = cv2.Canny(image, 100, 200)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    canny_image = Image.fromarray(image)
    return canny_image


def convert_to_grayscale(image):
    gray_image = image.convert('L').convert('RGB')
    return gray_image

def tile(downscale_factor, input_image):
    control_image = input_image.resize((input_image.size[0] // downscale_factor, input_image.size[1] // downscale_factor)).resize(input_image.size, Image.NEAREST)
    return control_image
    
def resize_img(control_image):
    image_ratio = control_image.width / control_image.height
    ratio = min(RATIO_CONFIGS_1024.keys(), key=lambda k: abs(k - image_ratio))
    to_height = RATIO_CONFIGS_1024[ratio]["height"]
    to_width = RATIO_CONFIGS_1024[ratio]["width"]
    resized_image = control_image.resize((to_width, to_height), resample=Image.Resampling.LANCZOS)
    return resized_image

@spaces.GPU(duration=180)
def infer(image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed, progress=gr.Progress(track_tqdm=True)):
    control_mode_num = mode_mapping[control_mode]
    
    if image_in is not None:
        image_in = resize_img(load_image(image_in))
        if control_mode == "canny":
            control_image = extract_canny(image_in)
        elif control_mode == "depth":
            control_image = extract_depth(image_in)
        elif control_mode == "pose":
            control_image = extract_openpose(image_in)
        elif control_mode == "colorgrid":
            control_image = tile(64, image_in)
        elif control_mode == "recolor":
            control_image = convert_to_grayscale(image_in)
        elif control_mode == "tile":
            control_image = tile(16, image_in)

    control_image = resize_img(control_image)

    width, height = control_image.size
    
    image = pipe(
        prompt, 
        control_image=control_image,
        control_mode=control_mode_num,
        width=width,
        height=height,
        controlnet_conditioning_scale=control_strength,
        num_inference_steps=inference_steps, 
        guidance_scale=guidance_scale,
        generator=torch.manual_seed(seed),
        max_sequence_length=128,
        negative_prompt="Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate"
    ).images[0]

    torch.cuda.empty_cache() 
    
    return image, control_image, gr.update(visible=True)
   

css="""
#col-container{
    margin: 0 auto;
    max-width: 1080px;
}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # BRIA-3.1-ControlNet-Union
        A unified ControlNet for BRIA-3.1 model from Bria.ai.<br />
        """)
        
        with gr.Column():
            
            with gr.Row():
                with gr.Column():
                    
                    # with gr.Row(equal_height=True):
                        # cond_in = gr.Image(label="Upload a processed control image", sources=["upload"], type="filepath")
                    image_in = gr.Image(label="Extract condition from a reference image (Optional)", sources=["upload"], type="filepath")
                    
                    prompt = gr.Textbox(label="Prompt", value="best quality")
                    
                    with gr.Accordion("Controlnet"):
                        control_mode = gr.Radio(
                            ["depth", "canny", "colorgrid", "recolor", "tile", "pose"], label="Mode", value="canny",
                            info="select the control mode, one for all"
                        )
                        
                        control_strength = gr.Slider(
                            label="control strength",
                            minimum=0,
                            maximum=1.0,
                            step=0.05,
                            value=0.9,
                        )
                    
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=555,
                    )
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    
                    with gr.Accordion("Advanced settings", open=False):
                        with gr.Column():
                            with gr.Row():
                                inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=50)
                                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0)
                    
                    submit_btn = gr.Button("Submit")
                    
                with gr.Column():
                    result = gr.Image(label="Result")
                    processed_cond = gr.Image(label="Preprocessed Cond")

    submit_btn.click(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False
    ).then(
        fn = infer,
        inputs = [image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed],
        outputs = [result, processed_cond],
        show_api=False
    )

demo.queue(api_open=False)
demo.launch()