app.py

import gradio as gr
import os
import torch
from diffusers import CogVideoXImageToVideoPipeline
from diffusers.utils import export_to_video, load_image
from datetime import datetime

from huggingface_hub import hf_hub_download

# Ensure 'checkpoint' directory exists
os.makedirs("checkpoints", exist_ok=True)

hf_hub_download(
    repo_id="wenqsun/DimensionX",
    filename="orbit_left_lora_weights.safetensors",
    local_dir="checkpoints"
)

hf_hub_download(
    repo_id="wenqsun/DimensionX",
    filename="orbit_up_lora_weights.safetensors",
    local_dir="checkpoints"
)

pipe = CogVideoXImageToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V", torch_dtype=torch.bfloat16)

def infer(prompt, image_path, orbit_type):
    lora_path = None
    if orbit_type == "Left":
        lora_path = "checkpoints/orbit_left_lora_weights.safetensors"
    elif orbit_type == "Up":
        lora_path = "checkpoints/orbit_up_lora_weights.safetensors"
    lora_rank = 256
    pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1")
    pipe.fuse_lora(lora_scale=1 / lora_rank)
    pipe.to("cuda")


    prompt = f"A{prompt}. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
    image = load_image(image_path)
    seed = random.randint(0, 2**8 - 1)
    
    video = pipe(
        image, 
        prompt, 
        num_inference_steps=50,  # NOT Changed
        guidance_scale=7.0,  # NOT Changed
        use_dynamic_cfg=True,
        generator=torch.Generator(device="cpu").manual_seed(seed)
    )
    
    # Generate a timestamp for the output filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    export_to_video(video.frames[0], f"output_{timestamp}.mp4", fps=8)
    return f"output_{timestamp}.mp4"

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# DimensionX")
        gr.Markdown("### Create Any 3D and 4D Scenes from a Single Image with Controllable Video Diffusion")
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="Image Input", type="filepath")
                prompt = gr.Textbox(label="Prompt")
                orbit_type = gr.Radio(label="Orbit type", choices=["Left", "Up"], value="Left")
                submit_btn = gr.Button("Submit")
            with gr.Column():
                video_out = gr.Video(label="Video output")
                examples = gr.Examples(
                    examples = [
                        [
                            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
                            "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background.",
                            "Left"
                        ]
                    ],
                    inputs=[image_in, prompt, orbit_type]
                )

    submit_btn.click(
        fn=infer,
        inputs=[image_in, prompt, orbit_type],
        outputs=[video_out]
    )

demo.queue().launch(show_error=True, show_api=False)