# built-in

import os
import subprocess
import logging
import re
import random
from string import ascii_letters, digits
import requests
import sys
import warnings

# external

#import spaces
import torch
import gradio as gr
from numpy import asarray as array
from lxml.html import fromstring
#from transformers import pipeline
#from diffusers.pipelines.flux import FluxPipeline
from diffusers.utils import export_to_gif, load_image
from diffusers.models.modeling_utils import ModelMixin
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file, save_file
from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, DDIMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, UNet3DConditionModel
#import jax
#import jax.numpy as jnp
from numba import cuda, njit as cpu, void, int64 as int, float64 as float, boolean as bool, uint8 as rgb
from numba.cuda import jit as gpu, grid, as_cuda_array as tensor2array
from numba.types import unicode_type as string
from PIL.Image import fromarray as array2image
import numpy as np

# logging

warnings.filterwarnings("ignore")
root = logging.getLogger()
root.setLevel(logging.DEBUG)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.DEBUG)
formatter = logging.Formatter('\n >>> [%(levelname)s] %(asctime)s %(name)s: %(message)s\n')
handler.setFormatter(formatter)
root.addHandler(handler)
handler2 = logging.StreamHandler(sys.stderr)
handler2.setLevel(logging.DEBUG)
formatter = logging.Formatter('\n >>> [%(levelname)s] %(asctime)s %(name)s: %(message)s\n')
handler2.setFormatter(formatter)
root.addHandler(handler2)

# output data

out_pipe=array([""])
last_motion=array([""])

# constant data

dtype = torch.float16
device = "cuda"
#repo = "ByteDance/AnimateDiff-Lightning"
#ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
base = "emilianJR/epiCRealism"
#base = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device, dtype=dtype)
#unet = UNet2DConditionModel.from_config("emilianJR/epiCRealism",subfolder="unet").to(device, dtype).load_state_dict(load_file(hf_hub_download("emilianJR/epiCRealism", "unet/diffusion_pytorch_model.safetensors"), device=device), strict=False)
adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-3", torch_dtype=dtype, device=device)

# precision data

fast=True
fps=10
time=1
width=384
height=768
step=40
accu=10

# ui data

css="""
input, input::placeholder {
    text-align: center !important;
}
*, *::placeholder {
    font-family: Suez One !important;
}
h1,h2,h3,h4,h5,h6 {
    width: 100%;
    text-align: center;
}
footer {
    display: none !important;
}
#col-container {
    margin: 0 auto;
    max-width: 15cm;
}
.image-container {
    aspect-ratio: """+str(width)+"/"+str(height)+""" !important;
}
.dropdown-arrow {
    display: none !important;
}
*:has(>.btn) {
    display: flex;
    justify-content: space-evenly;
    align-items: center;
}
.btn {
    display: flex;
}
"""

js="""
function custom(){
    document.querySelector("div#prompt input").setAttribute("maxlength","38")
    document.querySelector("div#prompt2 input").setAttribute("maxlength","38")
}
"""

# torch pipe

pipe = AnimateDiffPipeline.from_pretrained(base, vae=vae, motion_adapter=adapter, torch_dtype=dtype).to(device)
pipe.scheduler = DDIMScheduler(
    clip_sample=False,
    beta_start=0.00085,
    beta_end=0.012,
    beta_schedule="linear",
    timestep_spacing="trailing",
    steps_offset=1
)
pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
pipe.enable_vae_slicing()
pipe.enable_free_init(method="butterworth", use_fast_sampling=fast)


# functionality

def run(cmd):
    return str(subprocess.run(cmd, shell=True, capture_output=True, env=None).stdout)

def xpath_finder(str,pattern):
    try:
        return ""+fromstring(str).xpath(pattern)[0].text_content().lower().strip()
    except:
        return ""

def translate(text,lang):
    if text == None or lang == None:
        return ""       
    text = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', text)).lower().strip()
    lang = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', lang)).lower().strip()    
    if text == "" or lang == "":
        return ""
    if len(text) > 38:
        raise Exception("Translation Error: Too long text!")
    user_agents = [
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 13_1) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15'
    ]
    padded_chars = re.sub("[(^\-)(\-$)]","",text.replace("","-").replace("- -"," ")).strip()
    query_text = f'Please translate {padded_chars}, into {lang}'
    url = f'https://www.google.com/search?q={query_text}'
    content = str(requests.get(
        url = url,
        headers = {
            'User-Agent': random.choice(user_agents)
        }
    ).content)
    translated = text
    src_lang = xpath_finder(content,'//*[@class="source-language"]')
    trgt_lang = xpath_finder(content,'//*[@class="target-language"]')
    src_text = xpath_finder(content,'//*[@id="tw-source-text"]/*')
    trgt_text = xpath_finder(content,'//*[@id="tw-target-text"]/*')
    if trgt_lang == lang:
        translated = trgt_text
    ret = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', translated)).lower().strip()
    print(ret)
    return ret
    
def generate_random_string(length):
    characters = str(ascii_letters + digits)
    return ''.join(random.choice(characters) for _ in range(length))
    
@gpu(void( rgb[:], string[:], string[:], string[:] ))
def calc(img,p1,p2,motion):
    global out_pipe
    global last_motion
    global pipe

    x = grid(1)

    if last_motion[0] != motion:
        pipe.unload_lora_weights()
        if inp[3] != "":
            pipe.load_lora_weights(motion, adapter_name="motion")
            pipe.fuse_lora()
            pipe.set_adapters("motion", [0.7])
        last_motion[0] = motion

    pipe.to(device,dtype)

    if p2=="":
        out_pipe[x] = pipe(
            prompt=p1,
            height=height,
            width=width,
            ip_adapter_image=array2image(img).convert("RGB").resize((width,height)),
            num_inference_steps=step,
            guidance_scale=accu,
            num_frames=(fps*time)
        )
    
    out_pipe[x] = pipe(
        prompt=p1,
        negative_prompt=p2,
        height=height,
        width=width,
        ip_adapter_image=array2image(img).convert("RGB").resize((width,height)),
        num_inference_steps=step,
        guidance_scale=accu,
        num_frames=(fps*time)
    )

def handle(*inp):
    
    inp[1] = translate(inp[1],"english")
    inp[2] = translate(inp[2],"english")

    if inp[0] == None:
        return None

    if inp[2] != "":
        inp[2] = f"{inp[2]} where in the image"

    _do = ['photographed', 'realistic', 'dynamic poze', 'deep field', 'reasonable', "natural", 'rough', 'best quality', 'focused', "highly detailed"]
    if inp[1] != "":
        _do.append(f"a new {inp[1]} content in the image")
    inp[1] = ", ".join(_do)
    
    ln = len(result)

    inp[0] = array(inp[0])
    inp[1] = array(inp[1])
    inp[2] = array(inp[2])
    inp[3] = array(inp[3])
    calc[ln,32](*inp)

    for i in range(ln):
        name = generate_random_string(12)+".png"
        export_to_gif(out_pipe[i].frames[0],name,fps=fps)
        out_pipe[i] = name
    
    return out_pipe

def ui():
    with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
        with gr.Column(elem_id="col-container"):
            gr.Markdown(f"""
                # MULTI-LANGUAGE GIF CREATOR
            """)
            with gr.Row():
                global img
                img = gr.Image(label="STATIC PHOTO",show_label=True,container=True,type="pil")
            with gr.Row():
                global prompt
                prompt = gr.Textbox(
                    elem_id="prompt",
                    placeholder="INCLUDE",
                    container=False,
                    max_lines=1
                )
            with gr.Row():
                global prompt2
                prompt2 = gr.Textbox(
                    elem_id="prompt2",
                    placeholder="EXCLUDE",
                    container=False,
                    max_lines=1
                )
            with gr.Row():
                    global motion
                    motion = gr.Dropdown(
                        label='CAMERA',
                        show_label=True,
                        container=True,
                        choices=[
                            ("(No Effect)", ""),
                            ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),
                            ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),
                            ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),
                            ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),
                            ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),
                            ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),
                            ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),
                            ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),
                        ],
                        value="",
                        interactive=True
                    )
            with gr.Row():
                global run_button
                run_button = gr.Button("START",elem_classes="btn",scale=0)
            with gr.Row():
                global result
                result = []
                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
        demo.queue().launch()

@gpu(void())
def events():
    gr.on(
        triggers=[
            run_button.click,
            prompt.submit,
            prompt2.submit
        ],
        fn=handle,
        inputs=[img,prompt,prompt2,motion],
        outputs=result
    )

def entry():
    os.chdir(os.path.abspath(os.path.dirname(__file__)))
    
    ui()
    events[1,32]()

# entry

entry()

# end