Add "cpu" device choice

.python-version

@@ -0,0 +1 @@
+3.11.11

MotionDirector_inference.py

@@ -12,6 +12,7 @@
 from torch.nn.functional import interpolate
 from tqdm import trange
 import random
+import numpy as np
 
 from MotionDirector_train import export_to_video, handle_memory_attention, load_primary_models, unet_and_text_g_c, freeze_models
 from utils.lora_handler import LoraHandler
@@ -80,37 +81,114 @@ def inverse_video(pipe, latents, num_steps):
     return ddim_inv_latent
 
 
+# def prepare_input_latents(
+#     pipe: TextToVideoSDPipeline,
+#     batch_size: int,
+#     num_frames: int,
+#     height: int,
+#     width: int,
+#     latents_path:str,
+#     noise_prior: float,
+#     device: str = "cuda"
+# ):
+#     # initialize with random gaussian noise
+#     scale = pipe.vae_scale_factor
+#     shape = (batch_size, pipe.unet.config.in_channels, num_frames, height // scale, width // scale)
+#     print("shape: ", shape)
+#     if noise_prior > 0.:
+#         cached_latents = torch.load(latents_path, map_location=torch.device(device))
+#         for key in cached_latents:
+#             try:
+#                 print(f"cached_latents Key: {key}, Value:\n{cached_latents[key].shape}\n")
+#             except:
+#                 print(f"cached_latents Key: {key}, Value:\n{cached_latents[key]}\n")
+#         if 'inversion_noise' not in cached_latents:
+#             latents = inverse_video(pipe, cached_latents['latents'].unsqueeze(0), 50).squeeze(0)
+#             print("latents1.shape: ", latents.shape)
+#         else:
+#             latents = torch.load(latents_path)['inversion_noise'].unsqueeze(0)
+#             print("latents2.shape: ", latents.shape)
+#         if latents.shape[0] != batch_size:
+#             latents = latents.repeat(batch_size, 1, 1, 1, 1)
+#             print("latents3.shape: ", latents.shape)
+#         if latents.shape != shape:
+#             latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), (height // scale, width // scale), mode='bilinear')
+#             print("latents4.shape: ", latents.shape)
+#             latents = rearrange(latents, "(b f) c h w -> b c f h w", b=batch_size)
+#             print("latents5.shape: ", latents.shape)
+#         noise = torch.randn_like(latents, dtype=torch.half)
+#         print("noise.shape: ", noise.shape)
+#         latents = (noise_prior) ** 0.5 * latents + (1 - noise_prior) ** 0.5 * noise
+#         print("latents.shape: ", latents.shape)
+#     else:
+#         latents = torch.randn(shape, dtype=torch.half)
+
+#     print("latents.shape: ", latents.shape)
+
+#     return latents
+
 def prepare_input_latents(
     pipe: TextToVideoSDPipeline,
     batch_size: int,
     num_frames: int,
     height: int,
     width: int,
-    latents_path:str,
-    noise_prior: float
-):
-    # initialize with random gaussian noise
+    latents_path: str,
+    noise_prior: float,
+    device: str = "cuda"
+) -> Tensor:
+    # Initialize with random gaussian noise
     scale = pipe.vae_scale_factor
     shape = (batch_size, pipe.unet.config.in_channels, num_frames, height // scale, width // scale)
+    print("shape: ", shape)
+
     if noise_prior > 0.:
-        cached_latents = torch.load(latents_path)
+        cached_latents = torch.load(latents_path, map_location=torch.device(device))
+        for key in cached_latents:
+            try:
+                print(f"cached_latents Key: {key}, Value:\n{cached_latents[key].shape}\n")
+            except:
+                print(f"cached_latents Key: {key}, Value:\n{cached_latents[key]}\n")
+
         if 'inversion_noise' not in cached_latents:
             latents = inverse_video(pipe, cached_latents['latents'].unsqueeze(0), 50).squeeze(0)
+            print("latents1.shape: ", latents.shape)
         else:
             latents = torch.load(latents_path)['inversion_noise'].unsqueeze(0)
+            print("latents2.shape: ", latents.shape)
+
         if latents.shape[0] != batch_size:
             latents = latents.repeat(batch_size, 1, 1, 1, 1)
-        if latents.shape != shape:
-            latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), (height // scale, width // scale), mode='bilinear')
+            print("latents3.shape: ", latents.shape)
+
+        # Adjust num_frames if necessary
+        if latents.shape[2] != num_frames:
+            # Rearrange to (batch, channels, height, width, frames) for temporal interpolation
+            latents = rearrange(latents, "b c f h w -> b c h w f")
+            # Interpolate along the frame dimension
+            latents = interpolate(latents, size=(latents.shape[2], latents.shape[3], num_frames), mode='trilinear', align_corners=False)
+            # Rearrange back to (batch, channels, frames, height, width)
+            latents = rearrange(latents, "b c h w f -> b c f h w")
+            print("latents_temporal.shape: ", latents.shape)
+
+        if latents.shape[3:] != shape[3:]:
+            # Spatial interpolation
+            latents = interpolate(rearrange(latents, "b c f h w -> (b f) c h w", b=batch_size), 
+                                size=(height // scale, width // scale), mode='bilinear')
+            print("latents4.shape: ", latents.shape)
             latents = rearrange(latents, "(b f) c h w -> b c f h w", b=batch_size)
-        noise = torch.randn_like(latents, dtype=torch.half)
-        latents = (noise_prior) ** 0.5 * latents + (1 - noise_prior) ** 0.5 * noise
+            print("latents5.shape: ", latents.shape)
+
+        noise = torch.randn_like(latents, dtype=torch.float16)
+        print("noise.shape: ", noise.shape)
+        latents = (noise_prior ** 0.5) * latents + ((1 - noise_prior) ** 0.5) * noise
+        print("latents.shape: ", latents.shape)
     else:
-        latents = torch.randn(shape, dtype=torch.half)
-
+        latents = torch.randn(shape, dtype=torch.float16)
+        print("latents.shape: ", latents.shape)
+
     return latents
 
-
 def encode(pipe: TextToVideoSDPipeline, pixels: Tensor, batch_size: int = 8):
     nf = pixels.shape[2]
     pixels = rearrange(pixels, "b c f h w -> (b f) c h w")
@@ -158,7 +236,7 @@ def inference(
     with torch.autocast(device, dtype=torch.half):
         # prepare models
         pipe = initialize_pipeline(model, device, xformers, sdp, lora_path, lora_rank, lora_scale)
-
+        
         for i in range(repeat_num):
             if seed is None:
                 random_seed = random.randint(100, 10000000)
@@ -172,9 +250,10 @@ def inference(
                 height=height,
                 width=width,
                 latents_path=latents_path,
-                noise_prior=noise_prior
+                noise_prior=noise_prior,
+                device=device
             )
-
+            
             with torch.no_grad():
                 video_frames = pipe(
                     prompt=prompt,
@@ -186,7 +265,7 @@ def inference(
                     guidance_scale=guidance_scale,
                     latents=init_latents
                 ).frames
-
+                
             # =========================================
             # ========= write outputs to file =========
             # =========================================
@@ -229,7 +308,7 @@ def inference(
     parser.add_argument("-ls", "--lora_scale", type=float, default=1.0, help="Scale of LoRAs.")
     parser.add_argument("-r", "--seed", type=int, default=None, help="Random seed to make generations reproducible.")
     parser.add_argument("-np", "--noise_prior", type=float, default=0., help="Scale of the influence of inversion noise.")
-    parser.add_argument("-ci", "--checkpoint_index", type=int, required=True,
+    parser.add_argument("-ci", "--checkpoint_index", type=int, required=False, default=None,
                         help="The index of checkpoint, such as 300.")
     parser.add_argument("-rn", "--repeat_num", type=int, default=1,
                         help="How many results to generate with the same prompt.")
@@ -252,8 +331,11 @@ def inference(
     # =========================================
     # ============= sample videos =============
     # =========================================
-
-    lora_path = f"{args.checkpoint_folder}/checkpoint-{args.checkpoint_index}/temporal/lora"
+    if args.checkpoint_index:
+        lora_path = f"{args.checkpoint_folder}/checkpoint-{args.checkpoint_index}/temporal/lora"
+    else:
+        lora_path = f"{args.checkpoint_folder}/temporal/lora"
+
     latents_folder = f"{args.checkpoint_folder}/cached_latents"
     latents_path = f"{latents_folder}/{random.choice(os.listdir(latents_folder))}"
     assert os.path.exists(lora_path)
@@ -277,6 +359,4 @@ def inference(
         noise_prior=args.noise_prior,
         repeat_num=args.repeat_num
     )
-
-
-
+

MotionDirector_train.py

@@ -43,6 +43,25 @@
 import imageio
 import numpy as np
 
+try:
+    import torch_xla
+    import torch_xla.core.xla_model as xm
+    xla_available = True
+except ImportError:
+    xla_available = False
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+# Check for TPU (XLA device)
+elif xla_available:
+    DEVICE = xm.xla_device()
+# Fallback to CPU
+else:
+    DEVICE = "cpu"
+
+latents_device = "cpu" if DEVICE != "cuda" else "cuda"
+
+torch_dtype = torch.float16 if DEVICE == "cuda" else torch.float32
 
 already_printed_trainables = False
 
@@ -103,12 +122,36 @@ def extend_datasets(datasets, dataset_items, extend=False):
                     print(f"New {item} dataset length: {dataset.__len__()}")
                     extended.append(item)
 
+def export_to_video(video_frames, output_path, fps):
+    # Ensure video_frames is a list or tensor of shape (batch_size, num_frames, channels, height, width)
+    if isinstance(video_frames, torch.Tensor):
+        video_frames = video_frames.cpu().numpy()  # Convert tensor to NumPy
+    elif isinstance(video_frames, list):
+        video_frames = np.array(video_frames)
+
+    # Check shape and adjust if necessary
+    if len(video_frames.shape) == 5:  # (batch_size, num_frames, channels, height, width)
+        video_frames = video_frames[0]  # Take first batch if batch_size > 1
+    if video_frames.shape[1] == 3 or video_frames.shape[1] == 4:  # Channels in second dim (num_frames, channels, height, width)
+        video_frames = video_frames.transpose(0, 2, 3, 1)  # Reshape to (num_frames, height, width, channels)
+
+    # Ensure pixel values are in [0, 255] and uint8
+    if video_frames.max() <= 1.0:
+        video_frames = (video_frames * 255).astype(np.uint8)
+    else:
+        video_frames = video_frames.astype(np.uint8)
+
+    # Ensure exactly 3 channels (RGB)
+    if video_frames.shape[-1] == 4:  # If RGBA, drop alpha channel
+        video_frames = video_frames[..., :3]
+    elif video_frames.shape[-1] == 1:  # If grayscale, convert to RGB
+        video_frames = np.repeat(video_frames, 3, axis=-1)
 
-def export_to_video(video_frames, output_video_path, fps):
-    video_writer = imageio.get_writer(output_video_path, fps=fps)
-    for img in video_frames:
-        video_writer.append_data(np.array(img))
-    video_writer.close()
+    # Write video
+    writer = imageio.get_writer(output_path, fps=fps, codec='libx264')
+    for frame in video_frames:
+        writer.append_data(frame)
+    writer.close()
 
 
 def create_output_folders(output_dir, config):
@@ -131,11 +174,17 @@ def load_primary_models(pretrained_model_path):
 
     return noise_scheduler, tokenizer, text_encoder, vae, unet
 
-
 def unet_and_text_g_c(unet, text_encoder, unet_enable, text_enable):
-    unet._set_gradient_checkpointing(value=unet_enable)
-    text_encoder._set_gradient_checkpointing(CLIPEncoder, value=text_enable)
-
+    if hasattr(unet, '_set_gradient_checkpointing'):
+        print("unet._set_gradient_checkpointing(unet_enable)")
+        unet._set_gradient_checkpointing(unet_enable)
+    else:
+        print("NO unet._set_gradient_checkpointing(unet_enable)")
+
+    if hasattr(text_encoder, '_set_gradient_checkpointing'):
+        text_encoder._set_gradient_checkpointing(text_enable)
+    else:
+        print("NO text_encoder._set_gradient_checkpointing(text_enable)")
 
 def freeze_models(models_to_freeze):
     for model in models_to_freeze:
@@ -306,15 +355,15 @@ def handle_cache_latents(
     # Cache latents by storing them in VRAM.
     # Speeds up training and saves memory by not encoding during the train loop.
     if not should_cache: return None
-    vae.to('cuda', dtype=torch.float16)
+    vae.to(DEVICE, dtype=torch_dtype)
     vae.enable_slicing()
 
     pipe = TextToVideoSDPipeline.from_pretrained(
         pretrained_model_path,
         vae=vae,
-        unet=copy.deepcopy(unet).to('cuda', dtype=torch.float16)
+        unet=copy.deepcopy(unet).to(DEVICE, dtype=torch_dtype)
     )
-    pipe.text_encoder.to('cuda', dtype=torch.float16)
+    pipe.text_encoder.to(DEVICE, dtype=torch_dtype)
 
     cached_latent_dir = (
         os.path.abspath(cached_latent_dir) if cached_latent_dir is not None else None
@@ -329,7 +378,7 @@ def handle_cache_latents(
             save_name = f"cached_{i}"
             full_out_path = f"{cache_save_dir}/{save_name}.pt"
 
-            pixel_values = batch['pixel_values'].to('cuda', dtype=torch.float16)
+            pixel_values = batch['pixel_values'].to(DEVICE, dtype=torch_dtype)
             batch['latents'] = tensor_to_vae_latent(pixel_values, vae)
             if noise_prior > 0.:
                 batch['inversion_noise'] = inverse_video(pipe, batch['latents'], 50)
@@ -388,10 +437,10 @@ def tensor_to_vae_latent(t, vae):
 
 def sample_noise(latents, noise_strength, use_offset_noise=False):
     b, c, f, *_ = latents.shape
-    noise_latents = torch.randn_like(latents, device=latents.device)
+    noise_latents = torch.randn_like(latents, device=latents_device)
 
     if use_offset_noise:
-        offset_noise = torch.randn(b, c, f, 1, 1, device=latents.device)
+        offset_noise = torch.randn(b, c, f, 1, 1, device=latents_device)
         noise_latents = noise_latents + noise_strength * offset_noise
 
     return noise_latents
@@ -783,7 +832,7 @@ def finetune_unet(batch, step, mask_spatial_lora=False, mask_temporal_lora=False
         bsz = latents.shape[0]
 
         # Sample a random timestep for each video
-        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents_device)
         timesteps = timesteps.long()
 
         # Add noise to the latents according to the noise magnitude at each timestep

models/unet_3d_blocks.py

@@ -16,8 +16,8 @@
 import torch.utils.checkpoint as checkpoint
 from torch import nn
 from diffusers.models.resnet import Downsample2D, ResnetBlock2D, TemporalConvLayer, Upsample2D
-from diffusers.models.transformer_2d import Transformer2DModel
-from diffusers.models.transformer_temporal import TransformerTemporalModel
+from diffusers import Transformer2DModel
+from diffusers import TransformerTemporalModel
 
 # Assign gradient checkpoint function to simple variable for readability.
 g_c = checkpoint.checkpoint

models/unet_3d_condition.py

@@ -23,7 +23,7 @@
 from diffusers.utils import BaseOutput, logging
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.transformer_temporal import TransformerTemporalModel
+from diffusers import TransformerTemporalModel
 from .unet_3d_blocks import (
     CrossAttnDownBlock3D,
     CrossAttnUpBlock3D,