import random
import spaces
import torch
import gradio as gr

from modeling.dmm_pipeline import StableDiffusionDMMPipeline
from huggingface_hub import snapshot_download


ckpt_path = "ckpt"
snapshot_download(repo_id="MCG-NJU/DMM", local_dir=ckpt_path)

pipe = StableDiffusionDMMPipeline.from_pretrained(
    ckpt_path,
    torch_dtype=torch.float16, 
    use_safetensors=True
)
pipe.to("cuda")


@spaces.GPU
def generate(prompt: str,
             negative_prompt: str,
             model_id: int,
             seed: int = 1234,
             height: int = 512,
             width: int = 512,
             all: bool = True):
    if all:
        outputs = []
        for i in range(pipe.unet.get_num_models()):
            output = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=25,
                guidance_scale=7,
                model_id=i,
                generator=torch.Generator().manual_seed(seed),
            ).images[0]
            outputs.append(output)
        return outputs
    else:
        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            num_inference_steps=25,
            guidance_scale=7,
            model_id=int(model_id),
            generator=torch.Generator().manual_seed(seed),
        ).images[0]
        return [output,]


candidates = [
    "0. [JuggernautReborn] realistic",
    "1. [MajicmixRealisticV7] realistic, Asia portrait",
    "2. [EpicRealismV5] realistic",
    "3. [RealisticVisionV5] realistic",
    "4. [MajicmixFantasyV3] animation",
    "5. [MinimalismV2] illustration",
    "6. [RealCartoon3dV17] cartoon 3d",
    "7. [AWPaintingV1.4] animation",
]

def main():
    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # DMM Demo
            The checkpoint is https://huggingface.co/MCG-NJU/DMM.
            """
        )
        with gr.Row():
            with gr.Column():
                with gr.Column():
                    model_id = gr.Dropdown(candidates, label="Model Index", type="index")
                    all_check = gr.Checkbox(label="All (ignore the selection above)")
                prompt = gr.Textbox("portrait photo of a girl, long golden hair, flowers, best quality", label="Prompt")
                negative_prompt = gr.Textbox("worst quality,low quality,normal quality,lowres,watermark,nsfw", label="Negative Prompt")
                with gr.Row():
                    seed = gr.Number(0, label="Seed", precision=0, scale=3)
                    update_seed_btn = gr.Button("🎲", scale=1)
                with gr.Row():
                    height = gr.Number(768, step=8, label="Height (suggest 512~768)")
                    width = gr.Number(512, step=8, label="Width")
                submit_btn = gr.Button("Submit", variant="primary")
            output = gr.Gallery(label="images")
    
        submit_btn.click(generate,
                         inputs=[prompt, negative_prompt, model_id, seed, height, width, all_check],
                         outputs=[output])
        update_seed_btn.click(lambda: random.randint(0, 1000000),
                              outputs=[seed])

    demo.launch()


if __name__ == "__main__":
    main()