import gradio as gr
from PIL import Image
import torch
from diffusers import HunyuanVideoPipeline
# ... other imports ...

# Add LORA configuration
LORA_LIST = [
    "Top_Off.safetensors",
    "huanyan_helper.safetensors",
    "huanyan_helper_alpha.safetensors",
    "hunyuan-t-solo-v1.0.safetensors",
    "stripe_v2.safetensors"
]

def create_advanced_settings():
    with gr.Accordion("Advanced Settings", open=False):
        # LORA Selection
        lora_choices = gr.CheckboxGroup(
            choices=LORA_LIST,
            label="Select LORAs",
            value=[LORA_LIST[0]]
        )

        lora_weights = {}
        for lora in LORA_LIST:
            lora_weights[lora] = gr.Slider(0.0, 1.0, value=0.8,
                                        label=f"{lora} Weight")

        # Resolution Settings
        resolution = gr.Dropdown(
            choices=["512x512", "768x768", "1024x1024"],
            value="512x512",
            label="Output Resolution"
        )

    return lora_choices, lora_weights, resolution

def validate_image_resolution(image, resolution):
    if image is None:
        return
    img = Image.open(image)
    w, h = img.size
    if f"{w}x{h}" != resolution:
        raise gr.Error(f"Image resolution ({w}x{h}) must match output resolution ({resolution})")

def generate_video(prompt, negative_prompt, lora_choices, lora_weights,
                  resolution, image_input=None, steps=30):
    # Validate image resolution if provided
    if image_input:
        validate_image_resolution(image_input, resolution)

    # Load base model
    pipe = HunyuanVideoPipeline.from_pretrained(
        "Tencent-Hunyuan/Hunyuan-Video-Lite",
        torch_dtype=torch.float16
    ).to("cuda")

    # Apply selected LORAs
    for lora in lora_choices:
        pipe.load_lora_weights(
            f"TTV4ME/{lora}",
            adapter_name="hunyuanvideo-lora",
            weight_name=lora_weights[lora]
        )

    # Generate from image or text
    if image_input:
        image = Image.open(image_input).convert("RGB")
        output = pipe.image_to_video(
            image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_frames=24,
            height=int(resolution.split("x")[1]),
            width=int(resolution.split("x")[0]),
            num_inference_steps=steps
        )
    else:
        output = pipe.text_to_video(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=int(resolution.split("x")[1]),
            width=int(resolution.split("x")[0]),
            num_inference_steps=steps
        )

    return output.video

# Update interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative Prompt")
            image_input = gr.Image(label="Input Image", type="filepath")

            lora_choices, lora_weights, resolution = create_advanced_settings()

            generate_btn = gr.Button("Generate Video")

        with gr.Column():
            output_video = gr.Video(label="Generated Video")

    generate_btn.click(
        fn=generate_video,
        inputs=[prompt, negative_prompt, lora_choices,
               lora_weights, resolution, image_input],
        outputs=output_video
    )