from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch

# Load the model and tokenizer
model_name = "Skywork/SkyReels-V2-DF-1.3B-540P"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

def generate_video(prompt, max_length=512, temperature=0.7, top_k=50, top_p=0.95):
    """
    Generate video based on text prompt
    """
    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # Generate output
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            do_sample=True
        )
    
    # Decode and return the output
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # In a real implementation, you would process this into a video
    # For demo purposes, we'll just return the generated text
    return generated_text

# Create Gradio interface
iface = gr.Interface(
    fn=generate_video,
    inputs=gr.Textbox(lines=2, placeholder="Enter your video prompt here..."),
    outputs="text",
    title="SkyReels Video Generation",
    description="Generate video content using Skywork/SkyReels-V2-DF-1.3B-540P model",
    examples=[
        ["A sunny day at the beach with waves crashing"],
        ["A futuristic cityscape at night with flying cars"]
    ]
)

iface.launch()