import gradio as gr
import torch
from diffusers import LTXVideoPipeline
from huggingface_hub import login
import os

# Memuat pipeline LTX-Video dasar dari Hugging Face
# Model LTX-Video membutuhkan tipe data bfloat16 untuk efisiensi VRAM
pipe = LTXVideoPipeline.from_pretrained(
    "Lightricks/LTX-Video", 
    torch_dtype=torch.bfloat16
)

# Memuat bobot LoRA Black Magic milik FuzzPuppy ke dalam pipeline
pipe.load_lora_weights(
    "FuzzPuppy/LTX-2.3-Black-Magic-LoRA", 
    weight_name="pytorch_lora_weights.safetensors", 
    adapter_name="black_magic"
)
pipe.to("cuda")

def generate_video(prompt, negative_prompt, num_frames, fps, guidance_scale):
    # Mengonfigurasi parameter teks dan menjalankan inferensi
    video_frames = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=30,
        num_frames=int(num_frames),
        guidance_scale=float(guidance_scale),
        generator=torch.manual_seed(-1)
    ).frames[0]
    
    # Menyimpan frame video sementara ke format MP4
    output_path = "output_generated.mp4"
    # Catatan: Diffusers LTXVideo memproses penyimpanan video internal atau menggunakan manual export
    # Di bawah ini adalah contoh logika penyimpanan bawaan diffusers untuk video
    from diffusers.utils import export_to_video
    export_to_video(video_frames, output_path, fps=int(fps))
    
    return output_path

# Membuat Antarmuka Grafis Menggunakan Gradio
with gr.Blocks() as demo:
    gr.Markdown("# LTX-2.3 Black Magic LoRA Demo")
    gr.Markdown("Buat video cinematic magis menggunakan model dasar LTX-Video yang dipadukan dengan LoRA Black Magic.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="A wizard casting a dark magic spell, cinematic lighting, 4k...")
            negative_prompt = gr.Textbox(label="Negative Prompt", value="low quality, blurry, distorted")
            frames = gr.Slider(minimum=16, maximum=64, step=8, value=32, label="Jumlah Frame (Durasi)")
            fps = gr.Slider(minimum=8, maximum=24, step=2, value=16, label="FPS")
            guidance = gr.Slider(minimum=1.0, maximum=10.0, step=0.5, value=5.0, label="Guidance Scale")
            btn = gr.Button("Generate Video")
        
        with gr.Column():
            output_video = gr.Video(label="Hasil Video")
            
    btn.click(
        fn=generate_video, 
        inputs=[prompt, negative_prompt, frames, fps, guidance], 
        outputs=output_video
    )

demo.launch()