Merge pull request #9 from LightricksResearch/image-to-video-for-zeev

requirements.txt

@@ -3,4 +3,6 @@ diffusers==0.28.2
 transformers==4.44.2
 sentencepiece>=0.1.96
 accelerate
-einops
+einops
+matplotlib
+opencv-python

xora/examples/image_to_video.py

@@ -1,3 +1,4 @@
+import time
 import torch
 from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
 from xora.models.transformers.transformer3d import Transformer3DModel
@@ -9,6 +10,13 @@ from transformers import T5EncoderModel, T5Tokenizer
 import safetensors.torch
 import json
 import argparse
+from xora.utils.conditioning_method import ConditioningMethod
+import os
+import numpy as np
+import cv2
+from PIL import Image
+from tqdm import tqdm
+import random
 
 def load_vae(vae_dir):
     vae_ckpt_path = vae_dir / "diffusion_pytorch_model.safetensors"
@@ -34,78 +42,166 @@ def load_scheduler(scheduler_dir):
     scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
     return RectifiedFlowScheduler.from_config(scheduler_config)
 
+def center_crop_and_resize(frame, target_height, target_width):
+    h, w, _ = frame.shape
+    aspect_ratio_target = target_width / target_height
+    aspect_ratio_frame = w / h
+    if aspect_ratio_frame > aspect_ratio_target:
+        new_width = int(h * aspect_ratio_target)
+        x_start = (w - new_width) // 2
+        frame_cropped = frame[:, x_start:x_start + new_width]
+    else:
+        new_height = int(w / aspect_ratio_target)
+        y_start = (h - new_height) // 2
+        frame_cropped = frame[y_start:y_start + new_height, :]
+    frame_resized = cv2.resize(frame_cropped, (target_width, target_height))
+    return frame_resized
+
+def load_video_to_tensor_with_resize(video_path, target_height=512, target_width=768):
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame_resized = center_crop_and_resize(frame_rgb, target_height, target_width)
+        frames.append(frame_resized)
+    cap.release()
+    video_np = (np.array(frames) / 127.5) - 1.0
+    video_tensor = torch.tensor(video_np).permute(3, 0, 1, 2).float()
+    return video_tensor
+
+def load_image_to_tensor_with_resize(image_path, target_height=512, target_width=768):
+    image = Image.open(image_path).convert("RGB")
+    image_np = np.array(image)
+    frame_resized = center_crop_and_resize(image_np, target_height, target_width)
+    frame_tensor = torch.tensor(frame_resized).permute(2, 0, 1).float()
+    frame_tensor = (frame_tensor / 127.5) - 1.0
+    # Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)
+    return frame_tensor.unsqueeze(0).unsqueeze(2)
+
 def main():
-    # Parse command line arguments
-    parser = argparse.ArgumentParser(description='Load models from separate directories')
-    parser.add_argument('--separate_dir', type=str, required=True, help='Path to the directory containing unet, vae, and scheduler subdirectories')
+    parser = argparse.ArgumentParser(description='Load models from separate directories and run the pipeline.')
+
+    # Directories
+    parser.add_argument('--ckpt_dir', type=str, required=True,
+                        help='Path to the directory containing unet, vae, and scheduler subdirectories')
+    parser.add_argument('--video_path', type=str,
+                        help='Path to the input video file (first frame used)')
+    parser.add_argument('--image_path', type=str,
+                        help='Path to the input image file')
+    parser.add_argument('--seed', type=int, default="171198")
+
+    # Pipeline parameters
+    parser.add_argument('--num_inference_steps', type=int, default=40, help='Number of inference steps')
+    parser.add_argument('--num_images_per_prompt', type=int, default=1, help='Number of images per prompt')
+    parser.add_argument('--guidance_scale', type=float, default=3, help='Guidance scale for the pipeline')
+    parser.add_argument('--height', type=int, default=512, help='Height of the output video frames')
+    parser.add_argument('--width', type=int, default=768, help='Width of the output video frames')
+    parser.add_argument('--num_frames', type=int, default=121, help='Number of frames to generate in the output video')
+    parser.add_argument('--frame_rate', type=int, default=25, help='Frame rate for the output video')
+
+    # Prompts
+    parser.add_argument('--prompt', type=str,
+                        default='A man wearing a black leather jacket and blue jeans is riding a Harley Davidson motorcycle down a paved road. The man has short brown hair and is wearing a black helmet. The motorcycle is a dark red color with a large front fairing. The road is surrounded by green grass and trees. There is a gas station on the left side of the road with a red and white sign that says "Oil" and "Diner".',
+                        help='Text prompt to guide generation')
+    parser.add_argument('--negative_prompt', type=str,
+                        default='worst quality, inconsistent motion, blurry, jittery, distorted',
+                        help='Negative prompt for undesired features')
+
     args = parser.parse_args()
 
     # Paths for the separate mode directories
-    separate_dir = Path(args.separate_dir)
-    unet_dir = separate_dir / 'unet'
-    vae_dir = separate_dir / 'vae'
-    scheduler_dir = separate_dir / 'scheduler'
+    ckpt_dir = Path(args.ckpt_dir)
+    unet_dir = ckpt_dir / 'unet'
+    vae_dir = ckpt_dir / 'vae'
+    scheduler_dir = ckpt_dir / 'scheduler'
 
     # Load models
     vae = load_vae(vae_dir)
     unet = load_unet(unet_dir)
     scheduler = load_scheduler(scheduler_dir)
-
-    # Patchifier (remains the same)
     patchifier = SymmetricPatchifier(patch_size=1)
-
-    # text_encoder = T5EncoderModel.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder").to("cuda")
-    # tokenizer = T5Tokenizer.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer")
+    text_encoder = T5EncoderModel.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder").to(
+        "cuda")
+    tokenizer = T5Tokenizer.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer")
 
     # Use submodels for the pipeline
     submodel_dict = {
-        "transformer": unet,  # using unet for transformer
+        "transformer": unet,
         "patchifier": patchifier,
-        "text_encoder": None,
-        "tokenizer": None,
+        "text_encoder": text_encoder,
+        "tokenizer": tokenizer,
         "scheduler": scheduler,
         "vae": vae,
     }
 
-    model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
-    pipeline = VideoPixArtAlphaPipeline(
-                                                        **submodel_dict
-                                                        ).to("cuda")
-
-    num_inference_steps = 20
-    num_images_per_prompt = 1
-    guidance_scale = 3
-    height = 512
-    width = 768
-    num_frames = 57
-    frame_rate = 25
-
-    # Sample input stays the same
-    sample = torch.load("/opt/sample_media.pt")
-    for key, item in sample.items():
-        if item is not None:
-            sample[key] = item.cuda()
+    pipeline = VideoPixArtAlphaPipeline(**submodel_dict).to("cuda")
+
+    # Load media (video or image)
+    if args.video_path:
+        media_items = load_video_to_tensor_with_resize(args.video_path, args.height, args.width).unsqueeze(0)
+    elif args.image_path:
+        media_items = load_image_to_tensor_with_resize(args.image_path, args.height, args.width)
+    else:
+        raise ValueError("Either --video_path or --image_path must be provided.")
+
+    # Prepare input for the pipeline
+    sample = {
+        "prompt": args.prompt,
+        'prompt_attention_mask': None,
+        'negative_prompt': args.negative_prompt,
+        'negative_prompt_attention_mask': None,
+        'media_items': media_items,
+    }
 
-    # media_items = torch.load("/opt/sample_media.pt")
+    start_time = time.time()
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed(args.seed)
+    generator = torch.Generator(device="cuda").manual_seed(args.seed)
 
-    # Generate images (video frames)
     images = pipeline(
-        num_inference_steps=num_inference_steps,
-        num_images_per_prompt=num_images_per_prompt,
-        guidance_scale=guidance_scale,
-        generator=None,
+        num_inference_steps=args.num_inference_steps,
+        num_images_per_prompt=args.num_images_per_prompt,
+        guidance_scale=args.guidance_scale,
+        generator=generator,
         output_type="pt",
         callback_on_step_end=None,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        frame_rate=frame_rate,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
         **sample,
         is_video=True,
         vae_per_channel_normalize=True,
+        conditioning_method=ConditioningMethod.FIRST_FRAME
     ).images
+    # Save output video
+    def get_unique_filename(base, ext, dir='.', index_range=1000):
+        for i in range(index_range):
+            filename = os.path.join(dir, f"{base}_{i}{ext}")
+            if not os.path.exists(filename):
+                return filename
+        raise FileExistsError(f"Could not find a unique filename after {index_range} attempts.")
+
+
+    for i in range(images.shape[0]):
+        video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
+        video_np = (video_np * 255).astype(np.uint8)
+        fps = args.frame_rate
+        height, width = video_np.shape[1:3]
+        output_filename = get_unique_filename(f"video_output_{i}", ".mp4", ".")
+
+        out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+
+        for frame in video_np[..., ::-1]:
+            out.write(frame)
+
+        out.release()
 
-    print("Generated video frames.")
 
 if __name__ == "__main__":
     main()