import spaces
import gradio as gr
from sd3_pipeline import StableDiffusion3Pipeline
import torch
import random
import numpy as np
import os
import gc
from diffusers import AutoencoderKLWan
from wan_pipeline import WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from PIL import Image
from diffusers.utils import export_to_video
from huggingface_hub import login

# Authenticate with HF
login(token=os.getenv('HF_TOKEN'))

def set_seed(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)

# Updated model paths - now includes gated models
model_paths = {
    "sd2.1": "stabilityai/stable-diffusion-2-1",
    "sdxl": "stabilityai/stable-diffusion-xl-base-1.0",
    "sd3": "stabilityai/stable-diffusion-3-medium-diffusers",
    "sd3.5": "stabilityai/stable-diffusion-3.5-large",
    # "wan-t2v": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Keep commented if you don't have access to this one
}

current_model = None
OUTPUT_DIR = "generated_videos"
os.makedirs(OUTPUT_DIR, exist_ok=True)

def load_model(model_name):
    global current_model
    if current_model is not None:
        del current_model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

    # Determine device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    if "wan-t2v" in model_name:
        vae = AutoencoderKLWan.from_pretrained(
            model_paths[model_name], 
            subfolder="vae", 
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
        )
        scheduler = UniPCMultistepScheduler(
            prediction_type='flow_prediction', 
            use_flow_sigmas=True, 
            num_train_timesteps=1000, 
            flow_shift=8.0
        )
        current_model = WanPipeline.from_pretrained(
            model_paths[model_name], 
            vae=vae, 
            torch_dtype=torch.float16 if device == "cuda" else torch.float32
        ).to(device)
        current_model.scheduler = scheduler
    else:
        # Handle different model types
        if model_name in ["sd2.1"]:
            from diffusers import StableDiffusionPipeline
            current_model = StableDiffusionPipeline.from_pretrained(
                model_paths[model_name], 
                torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
            ).to(device)
        elif model_name in ["sdxl"]:
            from diffusers import StableDiffusionXLPipeline
            current_model = StableDiffusionXLPipeline.from_pretrained(
                model_paths[model_name], 
                torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
            ).to(device)
        else:
            # For SD3 models (when access is granted)
            current_model = StableDiffusion3Pipeline.from_pretrained(
                model_paths[model_name], 
                torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
            ).to(device)

    return current_model

@spaces.GPU(duration=120)
def generate_content(prompt, model_name, guidance_scale=7.5, num_inference_steps=50,
                     use_cfg_zero_star=True, use_zero_init=True, zero_steps=0,
                     seed=None, compare_mode=False):
    
    model = load_model(model_name)
    if seed is None:
        seed = random.randint(0, 2**32 - 1)
    set_seed(seed)

    is_video_model = "wan-t2v" in model_name
    print('prompt: ', prompt)
    
    if is_video_model:
        negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
        video1_frames = model(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=480,
            width=832,
            num_frames=81,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            use_cfg_zero_star=True,
            use_zero_init=True,
            zero_steps=zero_steps
        ).frames[0]
        video1_path = os.path.join(OUTPUT_DIR, f"{seed}_CFG-Zero-Star.mp4")
        export_to_video(video1_frames, video1_path, fps=16)

        return None, None, video1_path, seed

    # Handle different model types for image generation
    if model_name in ["sd2.1", "sdxl"]:
        # Standard diffusers pipeline interface
        if compare_mode:
            set_seed(seed)
            image1 = model(
                prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
            ).images[0]

            set_seed(seed)
            image2 = model(
                prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
            ).images[0]

            return image1, image2, None, seed
        else:
            image = model(
                prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
            ).images[0]

            return image, None, None, seed
    else:
        # SD3 models with custom parameters
        if compare_mode:
            set_seed(seed)
            image1 = model(
                prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                use_cfg_zero_star=True,
                use_zero_init=use_zero_init,
                zero_steps=zero_steps
            ).images[0]

            set_seed(seed)
            image2 = model(
                prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                use_cfg_zero_star=False,
                use_zero_init=use_zero_init,
                zero_steps=zero_steps
            ).images[0]

            return image1, image2, None, seed
        else:
            image = model(
                prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                use_cfg_zero_star=use_cfg_zero_star,
                use_zero_init=use_zero_init,
                zero_steps=zero_steps
            ).images[0]

            if use_cfg_zero_star:
                return image, None, None, seed
            else:
                return None, image, None, seed

# Gradio UI with left-right layout
with gr.Blocks() as demo:
    gr.HTML("""
        <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
            CFG-Zero*: Improved Classifier-Free Guidance for Flow Matching Models
        </div>
        <div style="text-align: center;">
            <a href="https://github.com/WeichenFan/CFG-Zero-star">Code</a> |
            <a href="https://arxiv.org/abs/2503.18886">Paper</a>
        </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(value="A spooky haunted mansion on a hill silhouetted by a full moon.", label="Enter your prompt")
            model_choice = gr.Dropdown(choices=list(model_paths.keys()), label="Choose Model")
            guidance_scale = gr.Slider(1, 20, value=4.0, step=0.5, label="Guidance Scale")
            inference_steps = gr.Slider(10, 100, value=50, step=5, label="Inference Steps")
            use_opt_scale = gr.Checkbox(value=True, label="Use Optimized-Scale")
            use_zero_init = gr.Checkbox(value=True, label="Use Zero Init")
            zero_steps = gr.Slider(0, 20, value=1, step=1, label="Zero out steps")
            seed = gr.Number(value=42, label="Seed (Leave blank for random)")
            compare_mode = gr.Checkbox(value=True, label="Compare Mode")
            generate_btn = gr.Button("Generate")

        with gr.Column(scale=2):
            out1 = gr.Image(type="pil", label="CFG-Zero* Image")
            out2 = gr.Image(type="pil", label="CFG Image")
            video = gr.Video(label="Video")
            used_seed = gr.Textbox(label="Used Seed")

    def update_params(model_name):
        print('model_name: ', model_name)
        if model_name == "wan-t2v":
            return (
                gr.update(value=5),
                gr.update(value=50),
                gr.update(value=True),
                gr.update(value=True),
                gr.update(value=1)
            )
        else:
            return (
                gr.update(value=4.0),
                gr.update(value=50),
                gr.update(value=True),
                gr.update(value=True),
                gr.update(value=1)
            )

    model_choice.change(
        fn=update_params,
        inputs=[model_choice],
        outputs=[guidance_scale, inference_steps, use_opt_scale, use_zero_init, zero_steps]
    )

    generate_btn.click(
        fn=generate_content,
        inputs=[
            prompt, model_choice, guidance_scale, inference_steps,
            use_opt_scale, use_zero_init, zero_steps, seed, compare_mode
        ],
        outputs=[out1, out2, video, used_seed]
    )

demo.launch(ssr_mode=False)