from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler, AutoencoderKL
import gradio as gr
import torch
from PIL import Image
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'


def convert_safetensors_to_bin(pipeline, state_dict, alpha = 0.4):
    LORA_PREFIX_UNET = 'lora_unet'
    LORA_PREFIX_TEXT_ENCODER = 'lora_te'

    visited = []

    # directly update weight in diffusers model
    for key in state_dict:
        # it is suggested to print out the key, it usually will be something like below
        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

        # as we have set the alpha beforehand, so just skip
        if '.alpha' in key or key in visited:
            continue

        if 'text' in key:
            layer_infos = key.split('.')[0].split(LORA_PREFIX_TEXT_ENCODER + '_')[-1].split('_')
            curr_layer = pipeline.text_encoder
        else:
            layer_infos = key.split('.')[0].split(LORA_PREFIX_UNET + '_')[-1].split('_')
            curr_layer = pipeline.unet

        # find the target layer
        temp_name = layer_infos.pop(0)
        while len(layer_infos) > -1:
            try:
                curr_layer = curr_layer.__getattr__(temp_name)
                if len(layer_infos) > 0:
                    temp_name = layer_infos.pop(0)
                elif len(layer_infos) == 0:
                    break
            except Exception:
                if len(temp_name) > 0:
                    temp_name += '_' + layer_infos.pop(0)
                else:
                    temp_name = layer_infos.pop(0)

        # org_forward(x) + lora_up(lora_down(x)) * multiplier
        pair_keys = []
        if 'lora_down' in key:
            pair_keys.append(key.replace('lora_down', 'lora_up'))
            pair_keys.append(key)
        else:
            pair_keys.append(key)
            pair_keys.append(key.replace('lora_up', 'lora_down'))

        # update weight
        if len(state_dict[pair_keys[0]].shape) == 4:
            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
            weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
        else:
            weight_up = state_dict[pair_keys[0]].to(torch.float32)
            weight_down = state_dict[pair_keys[1]].to(torch.float32)
            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)

        # update visited list
        for item in pair_keys:
            visited.append(item)

    return pipeline


model_id = 'andite/anything-v4.0'
prefix = ''
lora_path = hf_hub_download(
    "showee/showee-lora-v1.0", "showee-any4.0.safetensors"
)
vae_path = "./anything-v4.0-vae/diffusion_pytorch_model.bin"

scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")

pipe = StableDiffusionPipeline.from_pretrained(
  model_id,
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
  scheduler=scheduler)

pipe.vae.load_state_dict(torch.load(vae_path))

state_dict = load_file(lora_path)
pipe = convert_safetensors_to_bin(pipe, state_dict, 0.3)


pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
  model_id,
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
  scheduler=scheduler)

pipe_i2i.vae.load_state_dict(torch.load(vae_path))

state_dict_i2i = load_file(lora_path)
pipe_i2i = convert_safetensors_to_bin(pipe, state_dict_i2i, 0.3)


if torch.cuda.is_available():
  pipe = pipe.to("cuda")
  pipe_i2i = pipe_i2i.to("cuda")

def error_str(error, title="Error"):
    return f"""#### {title}
            {error}"""  if error else ""

def inference(prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt="", auto_prefix=False):
  if torch.cuda.is_available():
    generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
  else:
    generator = torch.Generator().manual_seed(seed) if seed != 0 else None
  prompt = f"{prefix} {prompt}" if auto_prefix else prompt
    
  try:
    if img is not None:
      return img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator), None
    else:
      return txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator), None
  except Exception as e:
    return None, error_str(e)

def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):

    result = pipe(
      prompt,
      negative_prompt = neg_prompt,
      num_inference_steps = int(steps),
      guidance_scale = guidance,
      width = width,
      height = height,
      generator = generator)
    
    return result.images[0]

def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):

    ratio = min(height / img.height, width / img.width)
    img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
    result = pipe_i2i(
        prompt,
        negative_prompt = neg_prompt,
        init_image = img,
        num_inference_steps = int(steps),
        strength = strength,
        guidance_scale = guidance,
        width = width,
        height = height,
        generator = generator)
        
    return result.images[0]

css = """.main-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.main-div div h1{font-weight:900;margin-bottom:7px}.main-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
"""
with gr.Blocks(css=css) as demo:
    gr.HTML(
        f"""
            <div class="main-div">
              <div>
                <h1>Showee V1.0</h1>
              </div>
              <p>
               Demo for <a href="https://huggingface.co/showee/showee-lora-v1.0">Showee V1.0</a> LoRA adaption weights fine-tuned from <a href="https://huggingface.co/andite/anything-v4.0">Anything V4.0</a> Stable Diffusion model.<br>
               {"Add the following tokens to your prompts for the model to work properly: <b>prefix</b>" if prefix else ""}
              </p>
              Running on {"<b>GPU 🔥</b>" if torch.cuda.is_available() else f"<b>CPU 🥶</b>. For faster inference it is recommended to <b>upgrade to GPU in <a href='https://huggingface.co/spaces/showee/showee-v1.0/settings'>Settings</a></b>"} after duplicating the space<br><br>
              <a style="display:inline-block" href="https://huggingface.co/spaces/showee/showee-v1.0?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
            </div>
        """
    )
    with gr.Row():
        
        with gr.Column(scale=55):
          with gr.Group():
              with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder=f"{prefix} [your prompt]").style(container=False)
                generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))

              image_out = gr.Image(height=512)
          error_output = gr.Markdown()

        with gr.Column(scale=45):
          with gr.Tab("Options"):
            with gr.Group():
              neg_prompt = gr.Textbox(label="Negative prompt",
                                      placeholder="What to exclude from the image",
                                      value="NSFW, lowres, ((bad anatomy)), ((bad hands)), text, missing finger, "
                                            "extra digits, fewer digits, blurry, ((mutated hands and fingers)), "
                                            "(poorly drawn face), ((mutation)), ((deformed face)), (ugly), "
                                            "((bad proportions)), ((extra limbs)), extra face, (double head), "
                                            "(extra head), ((extra feet)), monster, logo, cropped, worst quality, "
                                            "low quality, normal quality, jpeg, humpbacked, long body, long neck, "
                                            "((jpeg artifacts))")
              auto_prefix = gr.Checkbox(label="Prefix styling tokens automatically ()", value=prefix, visible=prefix)

              with gr.Row():
                guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)
                steps = gr.Slider(label="Steps", value=25, minimum=2, maximum=75, step=1)

              with gr.Row():
                width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8)
                height = gr.Slider(label="Height", value=512, minimum=64, maximum=1024, step=8)

              seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)

          with gr.Tab("Image to image"):
              with gr.Group():
                image = gr.Image(label="Image", height=256, tool="editor", type="pil")
                strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)

    gr.Examples(
        [[
            "masterpiece, best quality, ultra-detailed, illustration, portrait, 1girl, solo, white hair, green eyes, "
            "aqua_eyes, cat_ears, :3, ahoge, dress, red_jacket, long_sleeves, bangs, black_legwear, hair_ornament, "
            "hairclip",  8, 25, 768, 1024, 909198616
        ],
            [
                "masterpiece, best quality, ultra-detailed, illustration, portrait, 1girl, :3, animal_ears, aqua_eyes, ahoge, "
                "asymmetrical_legwear, bangs, black_footwear, black_skirt, breasts, cleavage, hair_ornament, hairclip, "
                "long_hair, navel, thighhighs, smile", 7.5, 25, 512, 768, 9
            ],
            [
                "masterpiece, best quality, ultra-detailed, illustration, portrait, 1girl, :3, animal_ears, aqua_eyes,  ahoge, seaside,"
                "asymmetrical_legwear, bangs, black_footwear, black_skirt, breasts, cleavage, hair_ornament, hairclip, "
                "long_hair, navel, thighhighs", 7.5, 25, 512, 512, 353573117
            ]],
        [prompt, guidance, steps, width, height, seed],
    )

    auto_prefix.change(lambda x: gr.update(placeholder=f"{prefix} [your prompt]" if x else "[Your prompt]"), inputs=auto_prefix, outputs=prompt, queue=False)

    inputs = [prompt, guidance, steps, width, height, seed, image, strength, neg_prompt, auto_prefix]
    outputs = [image_out, error_output]
    prompt.submit(inference, inputs=inputs, outputs=outputs)
    generate.click(inference, inputs=inputs, outputs=outputs)

    gr.HTML("""
    <div style="border-top: 1px solid #303030;">
      <br>
      <p>This space was created using <a href="https://huggingface.co/spaces/anzorq/sd-space-creator">SD Space Creator</a>.</p>
    </div>
    """)

demo.queue(concurrency_count=1)
demo.launch()