import torch
from diffusers import AutoencoderKLWan, WanPipeline, UniPCMultistepScheduler
from diffusers.utils import export_to_video
import gradio as gr
import tempfile
import spaces
from huggingface_hub import hf_hub_download
import numpy as np
import random
import os



from huggingface_hub import snapshot_download

snapshot_download(repo_id="APRIL-AIGC/UltraWan", repo_type="model", local_dir="ultrawan_weights/UltraWan", resume_download=True)



import subprocess
import os
import uuid

import subprocess

def upscale_to_4k(input_video_path, output_video_path):
    # Use Lanczos for better quality upscale
    cmd = [
        "ffmpeg",
        "-i", input_video_path,
        "-vf", "scale=3840:2160:flags=lanczos",  # upscale to 4K (3840x2160)
        "-c:v", "libx264",  # or libx265 for smaller size
        "-crf", "18",  # quality: lower is better (range 0-51)
        "-preset", "slow",  # better compression
        "-y",  # overwrite output file
        output_video_path,
    ]
    subprocess.run(cmd, check=True)




# LIGHT WEIGHT 1.3b
# MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
# LORA_REPO_ID = "Kijai/WanVideo_comfy"
# LORA_FILENAME = "Wan21_CausVid_bidirect2_T2V_1_3B_lora_rank32.safetensors"


MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"




LORA_REPO_ID = "Kijai/WanVideo_comfy"
LORA_FILENAME = "Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank256_bf16.safetensors"
#LORA_FILENAME = "Pusa/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors"




# LORA_REPO_ID = "RaphaelLiu/PusaV1"
# LORA_FILENAME="pusa_v1.safetensors"
#LORA_REPO_ID = "Kijai/WanVideo_comfy"
#LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"








vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
pipe = WanPipeline.from_pretrained(
    MODEL_ID, vae=vae, torch_dtype=torch.bfloat16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
pipe.to("cuda")

causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
pipe.set_adapters(["causvid_lora"], adapter_weights=[0.95])
pipe.fuse_lora()

# MOD_VALUE = 32
# DEFAULT_H_SLIDER_VALUE =  512
# DEFAULT_W_SLIDER_VALUE =  896

# # Environment variable check
# IS_ORIGINAL_SPACE = os.environ.get("IS_ORIGINAL_SPACE", "True") == "True"

# # Original limits
# ORIGINAL_SLIDER_MIN_H, ORIGINAL_SLIDER_MAX_H = 128, 1280
# ORIGINAL_SLIDER_MIN_W, ORIGINAL_SLIDER_MAX_W = 128, 1280
# ORIGINAL_MAX_DURATION = round(81/24, 1)  # MAX_FRAMES_MODEL/FIXED_FPS

# # Limited space constants
# LIMITED_MAX_RESOLUTION = 640
# LIMITED_MAX_DURATION = 2.0
# LIMITED_MAX_STEPS = 4

# # Set limits based on environment variable
# if IS_ORIGINAL_SPACE:
#     SLIDER_MIN_H, SLIDER_MAX_H = 128, LIMITED_MAX_RESOLUTION
#     SLIDER_MIN_W, SLIDER_MAX_W = 128, LIMITED_MAX_RESOLUTION
#     MAX_DURATION = LIMITED_MAX_DURATION
#     MAX_STEPS = LIMITED_MAX_STEPS
# else:
#     SLIDER_MIN_H, SLIDER_MAX_H = ORIGINAL_SLIDER_MIN_H, ORIGINAL_SLIDER_MAX_H
#     SLIDER_MIN_W, SLIDER_MAX_W = ORIGINAL_SLIDER_MIN_W, ORIGINAL_SLIDER_MAX_W
#     MAX_DURATION = ORIGINAL_MAX_DURATION
#     MAX_STEPS = 8

# MAX_SEED = np.iinfo(np.int32).max

# FIXED_FPS = 24
# FIXED_OUTPUT_FPS = 18 # we downspeed the output video as a temporary "trick"
# MIN_FRAMES_MODEL = 8
# MAX_FRAMES_MODEL = 81 


#New math to make it High Res

MOD_VALUE = 32

# Defaults for higher-res generation
DEFAULT_H_SLIDER_VALUE = 768
DEFAULT_W_SLIDER_VALUE = 1344  # 16:9 friendly and divisible by MOD_VALUE

# Original Space = Hugging Face space with compute limits
IS_ORIGINAL_SPACE = os.environ.get("IS_ORIGINAL_SPACE", "True") == "True"

# Conservative limits for low-end environments
LIMITED_MAX_RESOLUTION = 640
LIMITED_MAX_DURATION = 2.0
LIMITED_MAX_STEPS = 4

# Generous limits for local or Pro spaces
ORIGINAL_SLIDER_MIN_H, ORIGINAL_SLIDER_MAX_H = 128, 1536
ORIGINAL_SLIDER_MIN_W, ORIGINAL_SLIDER_MAX_W = 128, 1536
ORIGINAL_MAX_DURATION = round(81 / 24, 1)  # 3.4 seconds
ORIGINAL_MAX_STEPS = 8

# Use limited or original (generous) settings
if IS_ORIGINAL_SPACE:
    SLIDER_MIN_H, SLIDER_MAX_H = 128, LIMITED_MAX_RESOLUTION
    SLIDER_MIN_W, SLIDER_MAX_W = 128, LIMITED_MAX_RESOLUTION
    MAX_DURATION = LIMITED_MAX_DURATION
    MAX_STEPS = LIMITED_MAX_STEPS
else:
    SLIDER_MIN_H, SLIDER_MAX_H = ORIGINAL_SLIDER_MIN_H, ORIGINAL_SLIDER_MAX_H
    SLIDER_MIN_W, SLIDER_MAX_W = ORIGINAL_SLIDER_MIN_W, ORIGINAL_SLIDER_MAX_W
    MAX_DURATION = ORIGINAL_MAX_DURATION
    MAX_STEPS = ORIGINAL_MAX_STEPS

MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 24
FIXED_OUTPUT_FPS = 18  # reduce final video FPS to save space
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 81


default_prompt_t2v = "cinematic footage, group of pedestrians dancing in the streets of NYC, high quality breakdance, 4K, tiktok video, intricate details, instagram feel, dynamic camera, smooth dance motion, dimly lit, stylish, beautiful faces, smiling, music video"
default_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, watermark, text, signature"


import os
import tempfile
import random
import numpy as np
import torch
import gradio as gr
import subprocess
import shutil

def upscale_to_4k_and_replace(input_video_path):
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_upscaled:
        upscaled_path = tmp_upscaled.name

    cmd = [
        "ffmpeg",
        "-i", input_video_path,
        "-vf", "scale=3840:2160:flags=lanczos",
        "-c:v", "libx264",
        "-crf", "18",
        "-preset", "slow",
        "-y",
        upscaled_path,
    ]
    subprocess.run(cmd, check=True)

    shutil.move(upscaled_path, input_video_path)

def load_model_from_path(model_path: str):
    """
    Loads a diffusion pipeline from a local directory.
    The model is automatically loaded to CUDA with float16.
    """
    pipe = DiffusionPipeline.from_pretrained(
        model_path,
        torch_dtype=torch.float16,
        variant="fp16" if os.path.exists(os.path.join(model_path, "model.fp16.safetensors")) else None
    ).to("cuda")
    pipe.enable_model_cpu_offload()  # Optional: for large models
    return pipe

    
def get_duration(prompt, height, width, 
                   negative_prompt, duration_seconds,
                   guidance_scale, steps,
                   seed, randomize_seed,
                   progress):
    if steps > 4 and duration_seconds > 2:
        return 90
    elif steps > 4 or duration_seconds > 2:
        return 75
    else:
        return 60

@spaces.GPU(duration=get_duration)
def generate_video(prompt, height, width, 
                   negative_prompt=default_negative_prompt,
                   duration_seconds=2, guidance_scale=1,
                   steps=4, seed=42, randomize_seed=False,
                   progress=gr.Progress(track_tqdm=True)):

    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter a text prompt. Try to use long and precise descriptions.")

    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    # Clamp values in demo mode
    if IS_ORIGINAL_SPACE:
        height = min(height, LIMITED_MAX_RESOLUTION)
        width = min(width, LIMITED_MAX_RESOLUTION)
        duration_seconds = min(duration_seconds, LIMITED_MAX_DURATION)
        steps = min(steps, LIMITED_MAX_STEPS)

    # Ensure dimensions are aligned
    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)

    generator_pipe = pipe

    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)

    with torch.inference_mode():
        output_frames_list = generator_pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=target_h,
            width=target_w,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(steps),
            generator=torch.Generator(device="cuda").manual_seed(current_seed)
        ).frames[0]

    # Save video to temporary file
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    export_to_video(output_frames_list, video_path, fps=FIXED_OUTPUT_FPS)

    # Always upscale to 4K
    upscale_to_4k_and_replace(video_path)

    return video_path, current_seed


    
    

with gr.Blocks(css="body { max-width: 100vw; overflow-x: hidden; }") as demo:
    gr.HTML('<meta name="viewport" content="width=device-width, initial-scale=1">')
    # ... your other components here ...
    gr.Markdown("# ⚡ InstaVideo")
    gr.Markdown("This Gradio space is a fork of [wan2-1-fast from multimodalart](https://huggingface.co/spaces/multimodalart/wan2-1-fast), and is powered by the Wan CausVid LoRA [from Kijai](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_CausVid_bidirect2_T2V_1_3B_lora_rank32.safetensors).")

    # Add notice for limited spaces
    if IS_ORIGINAL_SPACE:
        gr.Markdown("⚠️ **This free public demo limits the resolution to 640px, duration to 2s, and inference steps to 4. For full capabilities please duplicate this space.**")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_t2v, placeholder="Describe the video you want to generate...")
            
            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
               
                with gr.Row():
                    height_input = gr.Slider(
                        minimum=SLIDER_MIN_H, 
                        maximum=SLIDER_MAX_H, 
                        step=MOD_VALUE, 
                        value=min(DEFAULT_H_SLIDER_VALUE, SLIDER_MAX_H), 
                        label=f"Output Height (multiple of {MOD_VALUE})"
                    )
                    width_input = gr.Slider(
                        minimum=SLIDER_MIN_W, 
                        maximum=SLIDER_MAX_W, 
                        step=MOD_VALUE, 
                        value=min(DEFAULT_W_SLIDER_VALUE, SLIDER_MAX_W), 
                        label=f"Output Width (multiple of {MOD_VALUE})"
                    )
                duration_seconds_input = gr.Slider(
                    minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), 
                    maximum=MAX_DURATION, 
                    step=0.1, 
                    value=2, 
                    label="Duration (seconds)", 
                    info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
                )
                steps_slider = gr.Slider(minimum=1, maximum=MAX_STEPS, step=1, value=4, label="Inference Steps") 
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)

            generate_button = gr.Button("Generate Video", variant="primary")
        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
    
    ui_inputs = [
        prompt_input, height_input, width_input,
        negative_prompt_input, duration_seconds_input,
        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
       
    ]
    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])

    # Adjust examples based on space limits
    example_configs = [
        ["a majestic eagle soaring through mountain peaks, cinematic aerial view", 896, 512],
        ["a serene ocean wave crashing on a sandy beach at sunset", 448, 832],
        ["a field of flowers swaying in the wind, spring morning light", 512, 896],
    ]
    
    if IS_ORIGINAL_SPACE:
        # Limit example resolutions for limited spaces
        example_configs = [
            [example[0], min(example[1], LIMITED_MAX_RESOLUTION), min(example[2], LIMITED_MAX_RESOLUTION)]
            for example in example_configs
        ]
    
    gr.Examples(
        examples=example_configs,
        inputs=[prompt_input, height_input, width_input], 
        outputs=[video_output, seed_input], 
        fn=generate_video, 
        cache_examples="lazy"
    )

if __name__ == "__main__":
    demo.queue().launch()