app.py

import torch
import gradio as gr
import imageio
import os
from safetensors.torch import load_file
from torchvision import transforms
from PIL import Image
import numpy as np

# Define model path (assuming it's in the HF Space)
MODEL_PATH = "sarthak247/Wan2.1-T2V-1.3B-nf4"
MODEL_FILE = f"{MODEL_PATH}/diffusion_pytorch_model.safetensors"

# Load model weights manually
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading model on {device}...")

try:
    model_weights = load_file(MODEL_FILE, device=device)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    model_weights = None

# Placeholder function - Replace with actual inference logic
def generate_video(prompt):
    """
    Generates a placeholder video using the model.
    Replace this function with the actual inference logic once available.
    """
    if model_weights is None:
        return "Model failed to load. Please check the logs."

    # Simulate an image output (Replace this with actual video frame generation)
    img = Image.new("RGB", (512, 512), color="black")
    transform = transforms.ToTensor()
    frame = (transform(img).permute(1, 2, 0).numpy() * 255).astype(np.uint8)

    # Create a fake video with repeated frames
    frames = [frame] * 16  # 16 repeated frames (Replace with actual video frames)
    output_path = "output.mp4"
    imageio.mimsave(output_path, frames, fps=8)

    return output_path

# Gradio UI
iface = gr.Interface(
    fn=generate_video,
    inputs=gr.Textbox(label="Enter Text Prompt"),
    outputs=gr.Video(label="Generated Video"),
    title="Wan2.1-T2V-1.3B Video Generation",
    description="This app loads the model manually and generates text-to-video output."
)

iface.launch()