################################################################################
# Copyright (C) 2023 Xingqian Xu - All Rights Reserved                         #
#                                                                              #
# Please visit Prompt-Free-Diffusion's arXiv paper for more details, link at   #
# arxiv.org/abs/2305.16223                                                     #
#                                                                              #
################################################################################

import gradio as gr
import os.path as osp
from PIL import Image
import numpy as np
import time

import torch
import torchvision.transforms as tvtrans
from lib.cfg_helper import model_cfg_bank
from lib.model_zoo import get_model

from collections import OrderedDict
from lib.model_zoo.ddim import DDIMSampler

n_sample_image = 1

controlnet_path = OrderedDict([
    ['canny'             , ('canny'   , 'pretrained/controlnet/control_sd15_canny_slimmed.safetensors')],
    ['canny_v11p'        , ('canny'   , 'pretrained/controlnet/control_v11p_sd15_canny_slimmed.safetensors')],
    ['depth'             , ('depth'   , 'pretrained/controlnet/control_sd15_depth_slimmed.safetensors')],
    ['hed'               , ('hed'     , 'pretrained/controlnet/control_sd15_hed_slimmed.safetensors')],
    ['mlsd'              , ('mlsd'    , 'pretrained/controlnet/control_sd15_mlsd_slimmed.safetensors')],
    ['mlsd_v11p'         , ('mlsd'    , 'pretrained/controlnet/control_v11p_sd15_mlsd_slimmed.safetensors')],
    ['normal'            , ('normal'  , 'pretrained/controlnet/control_sd15_normal_slimmed.safetensors')],
    ['openpose'          , ('openpose', 'pretrained/controlnet/control_sd15_openpose_slimmed.safetensors')],
    ['openpose_v11p'     , ('openpose', 'pretrained/controlnet/control_v11p_sd15_openpose_slimmed.safetensors')],
    ['scribble'          , ('scribble', 'pretrained/controlnet/control_sd15_scribble_slimmed.safetensors')],
    ['softedge_v11p'     , ('scribble', 'pretrained/controlnet/control_v11p_sd15_softedge_slimmed.safetensors')],
    ['seg'               , ('none'    , 'pretrained/controlnet/control_sd15_seg_slimmed.safetensors')],
    ['lineart_v11p'      , ('none'    , 'pretrained/controlnet/control_v11p_sd15_lineart_slimmed.safetensors')],
    ['lineart_anime_v11p', ('none'    , 'pretrained/controlnet/control_v11p_sd15s2_lineart_anime_slimmed.safetensors')],
])

preprocess_method = [
    'canny'                ,
    'depth'                ,
    'hed'                  ,
    'mlsd'                 ,
    'normal'               ,
    'openpose'             ,
    'openpose_withface'    ,
    'openpose_withfacehand',
    'scribble'             ,
    'none'                 ,
]

diffuser_path = OrderedDict([
    ['SD-v1.5'             , 'pretrained/pfd/diffuser/SD-v1-5.safetensors'],
    ['OpenJouney-v4'       , 'pretrained/pfd/diffuser/OpenJouney-v4.safetensors'],
    ['Deliberate-v2.0'     , 'pretrained/pfd/diffuser/Deliberate-v2-0.safetensors'],
    ['RealisticVision-v2.0', 'pretrained/pfd/diffuser/RealisticVision-v2-0.safetensors'],
    ['Anything-v4'         , 'pretrained/pfd/diffuser/Anything-v4.safetensors'],
    ['Oam-v3'              , 'pretrained/pfd/diffuser/AbyssOrangeMix-v3.safetensors'],
    ['Oam-v2'              , 'pretrained/pfd/diffuser/AbyssOrangeMix-v2.safetensors'],
])

ctxencoder_path = OrderedDict([
    ['SeeCoder'      , 'pretrained/pfd/seecoder/seecoder-v1-0.safetensors'],
    ['SeeCoder-PA'   , 'pretrained/pfd/seecoder/seecoder-pa-v1-0.safetensors'],
    ['SeeCoder-Anime', 'pretrained/pfd/seecoder/seecoder-anime-v1-0.safetensors'],
])

##########
# helper #
##########

def highlight_print(info):
    print('')
    print(''.join(['#']*(len(info)+4)))
    print('# '+info+' #')
    print(''.join(['#']*(len(info)+4)))
    print('')

def load_sd_from_file(target):
    if osp.splitext(target)[-1] == '.ckpt':
        sd = torch.load(target, map_location='cpu')['state_dict']
    elif osp.splitext(target)[-1] == '.pth':
        sd = torch.load(target, map_location='cpu')
    elif osp.splitext(target)[-1] == '.safetensors':
        from safetensors.torch import load_file as stload
        sd = OrderedDict(stload(target, device='cpu'))
    else:
        assert False, "File type must be .ckpt or .pth or .safetensors"
    return sd

########
# main #
########

class prompt_free_diffusion(object):
    def __init__(self, 
                 fp16=False, 
                 tag_ctx=None,
                 tag_diffuser=None,
                 tag_ctl=None,):

        self.tag_ctx = tag_ctx
        self.tag_diffuser = tag_diffuser
        self.tag_ctl = tag_ctl
        self.strict_sd = True

        cfgm = model_cfg_bank()('pfd_seecoder_with_controlnet')
        self.net = get_model()(cfgm)
        
        self.action_load_ctx(tag_ctx)
        self.action_load_diffuser(tag_diffuser)
        self.action_load_ctl(tag_ctl)
 
        if fp16:
            highlight_print('Running in FP16')
            self.net.ctx['image'].fp16 = True
            self.net = self.net.half()
            self.dtype = torch.float16
        else:
            self.dtype = torch.float32

        self.use_cuda = torch.cuda.is_available()
        if self.use_cuda:
            self.net.to('cuda')

        self.net.eval()
        self.sampler = DDIMSampler(self.net)

        self.n_sample_image = n_sample_image
        self.ddim_steps = 50
        self.ddim_eta = 0.0
        self.image_latent_dim = 4

    def load_ctx(self, pretrained):
        sd = load_sd_from_file(pretrained)
        sd_extra = [(ki, vi) for ki, vi in self.net.state_dict().items() \
            if ki.find('ctx.')!=0]
        sd.update(OrderedDict(sd_extra))

        self.net.load_state_dict(sd, strict=True)
        print('Load context encoder from [{}] strict [{}].'.format(pretrained, True))

    def load_diffuser(self, pretrained):
        sd = load_sd_from_file(pretrained)
        if len([ki for ki in sd.keys() if ki.find('diffuser.image.context_blocks.')==0]) == 0:
            sd = [(
                ki.replace('diffuser.text.context_blocks.', 'diffuser.image.context_blocks.'), vi) 
                    for ki, vi in sd.items()]
            sd = OrderedDict(sd)
        sd_extra = [(ki, vi) for ki, vi in self.net.state_dict().items() \
            if ki.find('diffuser.')!=0]
        sd.update(OrderedDict(sd_extra))
        self.net.load_state_dict(sd, strict=True)
        print('Load diffuser from [{}] strict [{}].'.format(pretrained, True))

    def load_ctl(self, pretrained):
        sd = load_sd_from_file(pretrained)
        self.net.ctl.load_state_dict(sd, strict=True)
        print('Load controlnet from [{}] strict [{}].'.format(pretrained, True))

    def action_load_ctx(self, tag):
        pretrained = ctxencoder_path[tag]
        if tag == 'SeeCoder-PA':
            from lib.model_zoo.seecoder import PPE_MLP
            pe_layer = \
                PPE_MLP(freq_num=20, freq_max=None, out_channel=768, mlp_layer=3)
            if self.dtype == torch.float16:
                pe_layer = pe_layer.half()
            if self.use_cuda:
                pe_layer.to('cuda')
            pe_layer.eval()
            self.net.ctx['image'].qtransformer.pe_layer = pe_layer
        else:
            self.net.ctx['image'].qtransformer.pe_layer = None
        if pretrained is not None:
            self.load_ctx(pretrained)
        self.tag_ctx = tag
        return tag

    def action_load_diffuser(self, tag):
        pretrained = diffuser_path[tag]
        if pretrained is not None:
            self.load_diffuser(pretrained)
        self.tag_diffuser = tag
        return tag

    def action_load_ctl(self, tag):
        pretrained = controlnet_path[tag][1]
        if pretrained is not None:
            self.load_ctl(pretrained)
        self.tag_ctl = tag
        return tag

    def action_autoset_hw(self, imctl):
        if imctl is None:
            return 512, 512
        w, h = imctl.size
        w = w//64 * 64
        h = h//64 * 64
        w = w if w >=512 else 512
        w = w if w <=1536 else 1536
        h = h if h >=512 else 512
        h = h if h <=1536 else 1536
        return h, w

    def action_autoset_method(self, tag):
        return controlnet_path[tag][0]

    def action_inference(
            self, im, imctl, ctl_method, do_preprocess, 
            h, w, ugscale, seed, 
            tag_ctx, tag_diffuser, tag_ctl,):

        if tag_ctx != self.tag_ctx:
            self.action_load_ctx(tag_ctx)
        if tag_diffuser != self.tag_diffuser:
            self.action_load_diffuser(tag_diffuser)
        if tag_ctl != self.tag_ctl:
            self.action_load_ctl(tag_ctl)

        n_samples = self.n_sample_image

        sampler = self.sampler
        device = self.net.device

        w = w//64 * 64
        h = h//64 * 64
        if imctl is not None:
            imctl = imctl.resize([w, h], Image.Resampling.BICUBIC)

        craw = tvtrans.ToTensor()(im)[None].to(device).to(self.dtype)
        c = self.net.ctx_encode(craw, which='image').repeat(n_samples, 1, 1)
        u = torch.zeros_like(c)

        if tag_ctx in ["SeeCoder-Anime"]:
            u = torch.load('assets/anime_ug.pth')[None].to(device).to(self.dtype)
            pad = c.size(1) - u.size(1)
            u = torch.cat([u, torch.zeros_like(u[:, 0:1].repeat(1, pad, 1))], axis=1)

        if tag_ctl != 'none':
            ccraw = tvtrans.ToTensor()(imctl)[None].to(device).to(self.dtype)
            if do_preprocess:
                cc = self.net.ctl.preprocess(ccraw, type=ctl_method, size=[h, w])
                cc = cc.to(self.dtype)
            else:
                cc = ccraw
        else:
            cc = None

        shape = [n_samples, self.image_latent_dim, h//8, w//8]

        if seed < 0:
            np.random.seed(int(time.time()))
            torch.manual_seed(-seed + 100)
        else:
            np.random.seed(seed + 100)
            torch.manual_seed(seed)

        x, _ = sampler.sample(
            steps=self.ddim_steps,
            x_info={'type':'image',},
            c_info={'type':'image', 'conditioning':c, 'unconditional_conditioning':u, 
                    'unconditional_guidance_scale':ugscale,
                    'control':cc,},
            shape=shape,
            verbose=False,
            eta=self.ddim_eta)

        ccout = [tvtrans.ToPILImage()(i) for i in cc] if cc is not None else []
        imout = self.net.vae_decode(x, which='image')
        imout = [tvtrans.ToPILImage()(i) for i in imout]
        return imout + ccout

pfd_inference = prompt_free_diffusion(
    fp16=True, tag_ctx = 'SeeCoder', tag_diffuser = 'Deliberate-v2.0', tag_ctl = 'canny',)

#################
# sub interface #
#################

cache_examples = True

def get_example():
    case = [
        [
            'assets/examples/ghibli-input.jpg', 
            'assets/examples/ghibli-canny.png', 
            'canny', False, 
            768, 1024, 1.8, 23, 
            'SeeCoder', 'Deliberate-v2.0', 'canny', ],
        [
            'assets/examples/astronautridinghouse-input.jpg', 
            'assets/examples/astronautridinghouse-canny.png', 
            'canny', False, 
            512, 768, 2.0, 21, 
            'SeeCoder', 'Deliberate-v2.0', 'canny', ],
        [
            'assets/examples/grassland-input.jpg', 
            'assets/examples/grassland-scribble.png', 
            'scribble', False, 
            768, 512, 2.0, 41, 
            'SeeCoder', 'Deliberate-v2.0', 'scribble', ],
        [
            'assets/examples/jeep-input.jpg', 
            'assets/examples/jeep-depth.png', 
            'depth', False, 
            512, 768, 2.0, 30, 
            'SeeCoder', 'Deliberate-v2.0', 'depth', ],
        [
            'assets/examples/bedroom-input.jpg', 
            'assets/examples/bedroom-mlsd.png', 
            'mlsd', False, 
            512, 512, 2.0, 31, 
            'SeeCoder', 'Deliberate-v2.0', 'mlsd', ],
        [
            'assets/examples/nightstreet-input.jpg', 
            'assets/examples/nightstreet-canny.png', 
            'canny', False, 
            768, 512, 2.3, 20, 
            'SeeCoder', 'Deliberate-v2.0', 'canny', ],
        [
            'assets/examples/woodcar-input.jpg', 
            'assets/examples/woodcar-depth.png', 
            'depth', False, 
            768, 512, 2.0, 20, 
            'SeeCoder', 'Deliberate-v2.0', 'depth', ],
        [
            'assets/examples-anime/miku.jpg', 
            'assets/examples-anime/miku-canny.png', 
            'canny', False, 
            768, 576, 1.5, 22, 
            'SeeCoder-Anime', 'Anything-v4', 'canny', ],
        [
            'assets/examples-anime/random0.jpg', 
            'assets/examples-anime/pose.png', 
            'openpose', False, 
            768, 1536, 2.0, 41, 
            'SeeCoder-Anime', 'Oam-v2', 'openpose_v11p', ],
        [
            'assets/examples-anime/random1.jpg', 
            'assets/examples-anime/pose.png', 
            'openpose', False, 
            768, 1536, 2.5, 28, 
            'SeeCoder-Anime', 'Oam-v2', 'openpose_v11p', ], 
        [
            'assets/examples-anime/camping.jpg', 
            'assets/examples-anime/pose.png', 
            'openpose', False, 
            768, 1536, 2.0, 35, 
            'SeeCoder-Anime', 'Anything-v4', 'openpose_v11p', ],
        [
            'assets/examples-anime/hanfu_girl.jpg', 
            'assets/examples-anime/pose.png', 
            'openpose', False, 
            768, 1536, 2.0, 20, 
            'SeeCoder-Anime', 'Anything-v4', 'openpose_v11p', ],
    ]
    return case

def interface():
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(label='Image Input', type='pil', elem_id='customized_imbox')
            with gr.Row():
                out_width  = gr.Slider(label="Width" , minimum=512, maximum=1536, value=512, step=64, visible=True)
                out_height = gr.Slider(label="Height", minimum=512, maximum=1536, value=512, step=64, visible=True)
            with gr.Row():
                scl_lvl = gr.Slider(label="CFGScale", minimum=0, maximum=10, value=2, step=0.01, visible=True)
                seed = gr.Number(20, label="Seed", precision=0)
            with gr.Row():
                tag_ctx = gr.Dropdown(label='Context Encoder', choices=[pi for pi in ctxencoder_path.keys()], value='SeeCoder')
                tag_diffuser = gr.Dropdown(label='Diffuser', choices=[pi for pi in diffuser_path.keys()], value='Deliberate-v2.0')
            button = gr.Button("Run")
        with gr.Column():
            ctl_input = gr.Image(label='Control Input', type='pil', elem_id='customized_imbox')
            do_preprocess = gr.Checkbox(label='Preprocess', value=False)
            with gr.Row():
                ctl_method = gr.Dropdown(label='Preprocess Type', choices=preprocess_method, value='canny')
                tag_ctl    = gr.Dropdown(label='ControlNet',      choices=[pi for pi in controlnet_path.keys()], value='canny')
        with gr.Column():
            img_output = gr.Gallery(label="Image Result", elem_id='customized_imbox').style(grid=n_sample_image+1)

    tag_ctl.change(
        pfd_inference.action_autoset_method,
        inputs = [tag_ctl],
        outputs = [ctl_method],)

    ctl_input.change(
        pfd_inference.action_autoset_hw,
        inputs = [ctl_input],
        outputs = [out_height, out_width],)

    # tag_ctx.change(
    #     pfd_inference.action_load_ctx,
    #     inputs = [tag_ctx],
    #     outputs = [tag_ctx],)

    # tag_diffuser.change(
    #     pfd_inference.action_load_diffuser,
    #     inputs = [tag_diffuser],
    #     outputs = [tag_diffuser],)

    # tag_ctl.change(
    #     pfd_inference.action_load_ctl,
    #     inputs = [tag_ctl],
    #     outputs = [tag_ctl],)

    button.click(
        pfd_inference.action_inference,
        inputs=[img_input, ctl_input, ctl_method, do_preprocess, 
                out_height, out_width, scl_lvl, seed, 
                tag_ctx, tag_diffuser, tag_ctl, ],
        outputs=[img_output])
    
    gr.Examples(
        label='Examples', 
        examples=get_example(), 
        fn=pfd_inference.action_inference,
        inputs=[img_input, ctl_input, ctl_method, do_preprocess,
                out_height, out_width, scl_lvl, seed, 
                tag_ctx, tag_diffuser, tag_ctl, ],
        outputs=[img_output],
        cache_examples=cache_examples,)

#############
# Interface #
#############

css = """
    #customized_imbox {
        min-height: 450px;
    }
    #customized_imbox>div[data-testid="image"] {
        min-height: 450px;
    }
    #customized_imbox>div[data-testid="image"]>div {
        min-height: 450px;
    }
    #customized_imbox>div[data-testid="image"]>iframe {
        min-height: 450px;
    }
    #customized_imbox>div.unpadded_box {
        min-height: 450px;
    }
    #myinst {
        font-size: 0.8rem; 
        margin: 0rem;
        color: #6B7280;
    }
    #maskinst {
        text-align: justify;
        min-width: 1200px;
    }
    #maskinst>img {
        min-width:399px;
        max-width:450px;
        vertical-align: top;
        display: inline-block;
    }
    #maskinst:after {
        content: "";
        width: 100%;
        display: inline-block;
    }
"""

if True:
    with gr.Blocks(css=css) as demo:
        gr.HTML(
            """
            <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
            <h1 style="font-weight: 900; font-size: 3rem; margin: 0rem">
                Prompt-Free Diffusion
            </h1>
            </div>
            """)

        interface()

        # gr.HTML(
        #     """
        #     <div style="text-align: justify; max-width: 1200px; margin: 20px auto;">
        #     <h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
        #     <b>Version</b>: {}
        #     </h3>
        #     </div>
        #     """.format(' '+str(pfd_inference.pretrained)))

    # demo.launch(server_name="0.0.0.0", server_port=7992)
    # demo.launch()
    demo.launch(debug=True)