Create app.py

app.py

@@ -0,0 +1,86 @@
+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)