import gradio as gr
import torch
import numpy as np
#import spaces
from utils import utils, tools, preprocess

#BASE_MODEL_REPO_ID = "neta-art/neta-xl-2.0"
#BASE_MODEL_FILENAME = "neta-xl-v2.fp16.safetensors"
BASE_MODEL_REPO_ID = "svjack/GenshinImpact_XL_Base"
BASE_MODEL_FILENAME = "sdxlBase_v10.safetensors"
VAE_PATH = "madebyollin/sdxl-vae-fp16-fix"
CONTROLNEXT_REPO_ID = "Eugeoter/controlnext-sdxl-anime-canny"
CACHE_DIR = None

DEFAULT_PROMPT = ""
DEFAULT_NEGATIVE_PROMPT = "worst quality, abstract, clumsy pose, deformed hand, dynamic malformation, fused fingers, extra digits, fewer digits, fewer fingers, extra fingers, extra arm, missing arm, extra leg, missing leg, signature, artist name, multi views, disfigured, ugly"


def ui():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipeline = tools.get_pipeline(
        pretrained_model_name_or_path=BASE_MODEL_REPO_ID,
        unet_model_name_or_path=CONTROLNEXT_REPO_ID,
        controlnet_model_name_or_path=CONTROLNEXT_REPO_ID,
        vae_model_name_or_path=VAE_PATH,
        load_weight_increasement=True,
        device=device,
        hf_cache_dir=CACHE_DIR,
        use_safetensors=True,
    )

    schedulers = ['Euler A', 'UniPC', 'Euler', 'DDIM', 'DDPM']

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

    with gr.Blocks(css=css) as demo:
        gr.Markdown(f"""
        # [ControlNeXt-SDXL](https://github.com/dvlab-research/ControlNeXt) Demo (Anime Canny)
        Base model: [Neta-Art-XL-2.0](https://civitai.com/models/410737/neta-art-xl)
        """)
        with gr.Row():
            with gr.Column(scale=9):
                prompt = gr.Textbox(label='Prompt', value=DEFAULT_PROMPT, lines=3, placeholder='prompt', container=False)
                negative_prompt = gr.Textbox(label='Negative Prompt', value=DEFAULT_NEGATIVE_PROMPT, lines=3, placeholder='negative prompt', container=False)
            with gr.Column(scale=1):
                generate_button = gr.Button("Generate", variant='primary', min_width=96)
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Row():
                    control_image = gr.Image(
                        value=None,
                        label='Condition',
                        sources=['upload'],
                        type='pil',
                        height=512,
                        image_mode='RGB',
                        format='png',
                        show_download_button=True,
                        show_share_button=True,
                    )
                with gr.Accordion(label='Preprocess', open=True):
                    with gr.Row():
                        threshold1 = gr.Slider(minimum=-1, maximum=255, step=1, value=-1, label='Threshold 1', info='-1 for auto')
                        threshold2 = gr.Slider(minimum=-1, maximum=255, step=1, value=-1, label='Threshold 2', info='-1 for auto')
                        process_button = gr.Button("Process", variant='primary', min_width=96, scale=0)
                with gr.Row():
                    scheduler = gr.Dropdown(
                        label='Scheduler',
                        choices=schedulers,
                        value='Euler A',
                        multiselect=False,
                        allow_custom_value=False,
                        filterable=True,
                    )
                    num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=28, label='Steps')
                with gr.Row():
                    cfg_scale = gr.Slider(minimum=1, maximum=30, step=1, value=7.5, label='CFG Scale')
                    controlnet_scale = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label='ControlNet Scale')
                with gr.Row():
                    seed = gr.Number(label='Seed', step=1, precision=0, value=-1)
            with gr.Column(scale=1):
                with gr.Row():
                    output = gr.Gallery(
                        label='Output',
                        value=None,
                        object_fit='scale-down',
                        columns=4,
                        height=512,
                        show_download_button=True,
                        show_share_button=True,
                    )

                with gr.Row():
                    examples = gr.Examples(
                        label='Examples',
                        examples=[
                            [
                                'best quality, lisa, solo, open hand, outdoors, indoor, cute, young, cat, cat ear, glasses',
                                'examples/example_1.jpg',
                            ],
                            [
                                'best quality, keqing, solo, the pose, indoors, smile',
                                'examples/example_2.jpg',
                            ]
                        ],
                        inputs=[
                            prompt,
                            control_image,
                        ],
                        cache_examples=False,
                    )

        #@spaces.GPU
        def generate(
            prompt,
            control_image,
            negative_prompt,
            cfg_scale,
            controlnet_scale,
            num_inference_steps,
            scheduler,
            seed,
        ):
            pipeline.scheduler = tools.get_scheduler(scheduler, pipeline.scheduler.config)
            generator = torch.Generator(device=device).manual_seed(max(0, min(seed, np.iinfo(np.int32).max))) if seed != -1 else None

            if control_image is None:
                raise gr.Error('Please upload an image.')
            width, height = utils.around_reso(control_image.width, control_image.height, reso=1024, max_width=2048, max_height=2048, divisible=32)
            control_image = control_image.resize((width, height)).convert('RGB')

            with torch.autocast(device):
                output_images = pipeline.__call__(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    controlnet_image=control_image,
                    controlnet_scale=controlnet_scale,
                    width=width,
                    height=height,
                    generator=generator,
                    guidance_scale=cfg_scale,
                    num_inference_steps=num_inference_steps,
                ).images

            return output_images

        def process(
            image,
            threshold1,
            threshold2,
        ):
            threshold1 = None if threshold1 == -1 else threshold1
            threshold2 = None if threshold2 == -1 else threshold2
            return preprocess.canny_extractor(image, threshold1, threshold2)

        generate_button.click(
            fn=generate,
            inputs=[prompt, control_image, negative_prompt, cfg_scale, controlnet_scale, num_inference_steps, scheduler, seed],
            outputs=[output],
        )

        process_button.click(
            fn=process,
            inputs=[control_image, threshold1, threshold2],
            outputs=[control_image],
        )

    return demo


if __name__ == '__main__':
    demo = ui()
    demo.queue().launch(share = True)