import os
import torch
import gradio as gr
from PIL import Image, ImageOps

from huggingface_hub import snapshot_download
from pyramid_dit import PyramidDiTForVideoGeneration
from diffusers.utils import export_to_video

import spaces 
import uuid

import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

# Constants
MODEL_PATH = "pyramid-flow-model"
MODEL_REPO = "rain1011/pyramid-flow-sd3"
MODEL_VARIANT = "diffusion_transformer_768p"
MODEL_DTYPE = "bf16"

def center_crop(image, target_width, target_height):
    width, height = image.size
    aspect_ratio_target = target_width / target_height
    aspect_ratio_image = width / height

    if aspect_ratio_image > aspect_ratio_target:
        # Crop the width (left and right)
        new_width = int(height * aspect_ratio_target)
        left = (width - new_width) // 2
        right = left + new_width
        top, bottom = 0, height
    else:
        # Crop the height (top and bottom)
        new_height = int(width / aspect_ratio_target)
        top = (height - new_height) // 2
        bottom = top + new_height
        left, right = 0, width

    image = image.crop((left, top, right, bottom))
    return image

# Download and load the model
def load_model():
    if not os.path.exists(MODEL_PATH):
        snapshot_download(MODEL_REPO, local_dir=MODEL_PATH, local_dir_use_symlinks=False, repo_type='model')
    
    model = PyramidDiTForVideoGeneration(
        MODEL_PATH,
        MODEL_DTYPE,
        model_variant=MODEL_VARIANT,
    )
    
    model.vae.to("cuda")
    model.dit.to("cuda")
    model.text_encoder.to("cuda")
    model.vae.enable_tiling()
    
    return model

# Global model variable
model = load_model()

# Text-to-video generation function
@spaces.GPU(duration=240)
def generate_video(image, prompt, duration, guidance_scale, video_guidance_scale):
    temp = int(duration * 0.8)  # Convert seconds to temp value (assuming 24 FPS)
    torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
    if(image):
        cropped_image = center_crop(image, 1280, 720)
        resized_image = cropped_image.resize((1280, 720))
        with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
            frames = model.generate_i2v(
                prompt=prompt,
                input_image=resized_image,
                num_inference_steps=[10, 10, 10],
                temp=temp,
                guidance_scale=7.0,
                video_guidance_scale=video_guidance_scale,
                output_type="pil",
                save_memory=True,
            )
    else:
        with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
            frames = model.generate(
                prompt=prompt,
                num_inference_steps=[20, 20, 20],
                video_num_inference_steps=[10, 10, 10],
                height=768,
                width=1280,
                temp=temp,
                guidance_scale=guidance_scale,
                video_guidance_scale=video_guidance_scale,
                output_type="pil",
                save_memory=True,
            )
    output_path = f"{str(uuid.uuid4())}_output_video.mp4"
    export_to_video(frames, output_path, fps=8)
    return output_path

# Image-to-video generation function
#@spaces.GPU(duration=240)
#def generate_video_from_image(image, prompt, duration, video_guidance_scale):
#    temp = int(duration * 2.4)  # Convert seconds to temp value (assuming 24 FPS)
#    torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
#    
#    target_size = (1280, 720)
#    cropped_image = center_crop(image, 1280, 720)
#    resized_image = cropped_image.resize((1280, 720))
#    
#    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
#        frames = model.generate_i2v(
#            prompt=prompt,
#            input_image=resized_image,
#            num_inference_steps=[10, 10, 10],
#            temp=temp,
#            guidance_scale=7.0,
#            video_guidance_scale=video_guidance_scale,
#            output_type="pil",
#            save_memory=True,
#        )
    
    output_path = "output_video_i2v.mp4"
    export_to_video(frames, output_path, fps=24)
    return output_path

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Pyramid Flow Video Generation Demo")
    
    #with gr.Tab("Text-to-Video"):
    with gr.Row():
        with gr.Column():
            with gr.Accordion("Image to Video (optional)", open=False):
                i2v_image = gr.Image(type="pil", label="Input Image")
            t2v_prompt = gr.Textbox(label="Prompt")
            with gr.Accordion("Advanced settings", open=False):
                t2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
                t2v_guidance_scale = gr.Slider(minimum=1, maximum=15, value=9, step=0.1, label="Guidance Scale")
                t2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=5, step=0.1, label="Video Guidance Scale")
            t2v_generate_btn = gr.Button("Generate Video")
        with gr.Column():
            t2v_output = gr.Video(label="Generated Video")
    
    t2v_generate_btn.click(
        generate_video,
        inputs=[i2v_image, t2v_prompt, t2v_duration, t2v_guidance_scale, t2v_video_guidance_scale],
        outputs=t2v_output
    )
    
    #with gr.Tab("Image-to-Video"):
    #    with gr.Row():
    #        with gr.Column():
                
    #            i2v_prompt = gr.Textbox(label="Prompt")
    #            i2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
    #            i2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=4, step=0.1, label="Video Guidance Scale")
    #            i2v_generate_btn = gr.Button("Generate Video")
    #        with gr.Column():
    #            i2v_output = gr.Video(label="Generated Video")
        
        #i2v_generate_btn.click(
        #    generate_video_from_image,
        #    inputs=[i2v_image, i2v_prompt, i2v_duration, i2v_video_guidance_scale],
        #    outputs=i2v_output
        #)

demo.launch()