Merge pull request #5 from LightricksResearch/safetensors-ckpts

scripts/to_safetensors.py

@@ -0,0 +1,109 @@
+import argparse
+from pathlib import Path
+from typing import Any, Dict
+import safetensors.torch
+import torch
+import json
+import shutil
+
+
+def load_text_encoder(index_path: Path) -> Dict:
+    with open(index_path, 'r') as f:
+        index: Dict = json.load(f)
+
+    loaded_tensors = {}
+    for part_file in set(index.get("weight_map", {}).values()):
+        tensors = safetensors.torch.load_file(index_path.parent / part_file, device='cpu')
+        for tensor_name in tensors:
+            loaded_tensors[tensor_name] = tensors[tensor_name]
+
+    return loaded_tensors
+
+
+def convert_unet(unet: Dict, add_prefix=True) -> Dict:
+    if add_prefix:
+        return {"model.diffusion_model." + key: value for key, value in unet.items()}
+    return unet
+
+
+def convert_vae(vae_path: Path, add_prefix=True) -> Dict:
+    state_dict = torch.load(vae_path / "autoencoder.pth", weights_only=True)
+    stats_path = vae_path / "per_channel_statistics.json"
+    if stats_path.exists():
+        with open(stats_path, 'r') as f:
+            data = json.load(f)
+        transposed_data = list(zip(*data["data"]))
+        data_dict = {
+            f"{'vae.' if add_prefix else ''}per_channel_statistics.{col}": torch.tensor(vals)
+            for col, vals in zip(data["columns"], transposed_data)
+        }
+    else:
+        data_dict = {}
+
+    result = {("vae." if add_prefix else "") + key: value for key, value in state_dict.items()}
+    result.update(data_dict)
+    return result
+
+
+def convert_encoder(encoder: Dict) -> Dict:
+    return {"text_encoders.t5xxl.transformer." + key: value for key, value in encoder.items()}
+
+
+def save_config(config_src: str, config_dst: str):
+    shutil.copy(config_src, config_dst)
+
+
+def load_vae_config(vae_path: Path) -> str:
+    config_path = vae_path / "config.json"
+    if not config_path.exists():
+        raise FileNotFoundError(f"VAE config file {config_path} not found.")
+    return str(config_path)
+
+
+def main(unet_path: str, vae_path: str, out_path: str, mode: str,
+         unet_config_path: str = None, scheduler_config_path: str = None) -> None:
+    unet = convert_unet(torch.load(unet_path, weights_only=True), add_prefix=(mode == 'single'))
+
+    # Load VAE from directory and config
+    vae = convert_vae(Path(vae_path), add_prefix=(mode == 'single'))
+    vae_config_path = load_vae_config(Path(vae_path))
+
+    if mode == 'single':
+        result = {**unet, **vae}
+        safetensors.torch.save_file(result, out_path)
+    elif mode == 'separate':
+        # Create directories for unet, vae, and scheduler
+        unet_dir = Path(out_path) / 'unet'
+        vae_dir = Path(out_path) / 'vae'
+        scheduler_dir = Path(out_path) / 'scheduler'
+
+        unet_dir.mkdir(parents=True, exist_ok=True)
+        vae_dir.mkdir(parents=True, exist_ok=True)
+        scheduler_dir.mkdir(parents=True, exist_ok=True)
+
+        # Save unet and vae safetensors with the name diffusion_pytorch_model.safetensors
+        safetensors.torch.save_file(unet, unet_dir / 'diffusion_pytorch_model.safetensors')
+        safetensors.torch.save_file(vae, vae_dir / 'diffusion_pytorch_model.safetensors')
+
+        # Save config files for unet, vae, and scheduler
+        if unet_config_path:
+            save_config(unet_config_path, unet_dir / 'config.json')
+        if vae_config_path:
+            save_config(vae_config_path, vae_dir / 'config.json')
+        if scheduler_config_path:
+            save_config(scheduler_config_path, scheduler_dir / 'scheduler_config.json')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--unet_path', '-u', type=str, default='unet/ema-002.pt')
+    parser.add_argument('--vae_path', '-v', type=str, default='vae/')
+    parser.add_argument('--out_path', '-o', type=str, default='xora.safetensors')
+    parser.add_argument('--mode', '-m', type=str, choices=['single', 'separate'], default='single',
+                        help="Choose 'single' for the original behavior, 'separate' to save unet and vae separately.")
+    parser.add_argument('--unet_config_path', type=str, help="Path to the UNet config file (for separate mode)")
+    parser.add_argument('--scheduler_config_path', type=str,
+                        help="Path to the Scheduler config file (for separate mode)")
+
+    args = parser.parse_args()
+    main(**args.__dict__)

xora/examples/image_to_video.py

@@ -5,83 +5,107 @@ from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier
 from xora.schedulers.rf import RectifiedFlowScheduler
 from xora.pipelines.pipeline_video_pixart_alpha import VideoPixArtAlphaPipeline
 from pathlib import Path
-from transformers import T5EncoderModel
-
-
-model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
-vae_local_path = Path("/opt/models/checkpoints/vae_training/causal_vvae_32x32x8_420m_cont_32/step_2296000")
-dtype = torch.float32
-vae = CausalVideoAutoencoder.from_pretrained(
-            pretrained_model_name_or_path=vae_local_path,
-            revision=False,
-            torch_dtype=torch.bfloat16,
-            load_in_8bit=False,
-).cuda()
-transformer_config_path = Path("/opt/txt2img/txt2img/config/transformer3d/xora_v1.2-L.json")
-transformer_config = Transformer3DModel.load_config(transformer_config_path)
-transformer = Transformer3DModel.from_config(transformer_config)
-transformer_local_path = Path("/opt/models/logs/v1.2-vae-mf-medHR-mr-cvae-first-frame-cond-4k-seq/ckpt/01822000/model.pt")
-transformer_ckpt_state_dict = torch.load(transformer_local_path)
-transformer.load_state_dict(transformer_ckpt_state_dict, True)
-transformer = transformer.cuda()
-unet = transformer
-scheduler_config_path = Path("/opt/txt2img/txt2img/config/scheduler/RF_SD3_shifted.json")
-scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
-scheduler = RectifiedFlowScheduler.from_config(scheduler_config)
-patchifier = SymmetricPatchifier(patch_size=1)
-# text_encoder = T5EncoderModel.from_pretrained("t5-v1_1-xxl")
-
-submodel_dict = {
-    "unet": unet,
-    "transformer": transformer,
-    "patchifier": patchifier,
-    "text_encoder": None,
-    "scheduler": scheduler,
-    "vae": vae,
-
-}
-
-pipeline = VideoPixArtAlphaPipeline.from_pretrained(model_name_or_path,
-                                                    safety_checker=None,
-            revision=None,
-            torch_dtype=dtype,
-            **submodel_dict,
-        )
-
-num_inference_steps=20
-num_images_per_prompt=2
-guidance_scale=3
-height=512
-width=768
-num_frames=57
-frame_rate=25
-# sample = {
-#     "prompt": "A cat", # (B, L, E)
-#     'prompt_attention_mask': None, # (B , L)
-#     'negative_prompt': "Ugly deformed",
-#     'negative_prompt_attention_mask': None # (B , L)
-# }
-
-sample = torch.load("/opt/sample.pt")
-for _, item in sample.items():
-    if item is not None:
-        item = item.cuda()
-media_items = torch.load("/opt/sample_media.pt")
-
-images = pipeline(
-    num_inference_steps=num_inference_steps,
-    num_images_per_prompt=num_images_per_prompt,
-    guidance_scale=guidance_scale,
-    generator=None,
-    output_type="pt",
-    callback_on_step_end=None,
-    height=height,
-    width=width,
-    num_frames=num_frames,
-    frame_rate=frame_rate,
-    **sample,
-    is_video=True,
-    vae_per_channel_normalize=True,
-).images
-
-print()
+from transformers import T5EncoderModel, T5Tokenizer
+import safetensors.torch
+import json
+import argparse
+
+def load_vae(vae_dir):
+    vae_ckpt_path = vae_dir / "diffusion_pytorch_model.safetensors"
+    vae_config_path = vae_dir / "config.json"
+    with open(vae_config_path, 'r') as f:
+        vae_config = json.load(f)
+    vae = CausalVideoAutoencoder.from_config(vae_config)
+    vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
+    vae.load_state_dict(vae_state_dict)
+    return vae.cuda().to(torch.bfloat16)
+
+def load_unet(unet_dir):
+    unet_ckpt_path = unet_dir / "diffusion_pytorch_model.safetensors"
+    unet_config_path = unet_dir / "config.json"
+    transformer_config = Transformer3DModel.load_config(unet_config_path)
+    transformer = Transformer3DModel.from_config(transformer_config)
+    unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
+    transformer.load_state_dict(unet_state_dict, strict=True)
+    return transformer.cuda()
+
+def load_scheduler(scheduler_dir):
+    scheduler_config_path = scheduler_dir / "scheduler_config.json"
+    scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
+    return RectifiedFlowScheduler.from_config(scheduler_config)
+
+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')
+    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'
+
+    # 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")
+
+    # Use submodels for the pipeline
+    submodel_dict = {
+        "transformer": unet,  # using unet for transformer
+        "patchifier": patchifier,
+        "text_encoder": None,
+        "tokenizer": None,
+        "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()
+
+    # media_items = torch.load("/opt/sample_media.pt")
+
+    # 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,
+        output_type="pt",
+        callback_on_step_end=None,
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        frame_rate=frame_rate,
+        **sample,
+        is_video=True,
+        vae_per_channel_normalize=True,
+    ).images
+
+    print("Generated video frames.")
+
+if __name__ == "__main__":
+    main()

xora/examples/text_to_video.py

@@ -5,84 +5,104 @@ from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier
 from xora.schedulers.rf import RectifiedFlowScheduler
 from xora.pipelines.pipeline_video_pixart_alpha import VideoPixArtAlphaPipeline
 from pathlib import Path
-from transformers import T5EncoderModel
+from transformers import T5EncoderModel, T5Tokenizer
+import safetensors.torch
+import json
+import argparse
 
+def load_vae(vae_dir):
+    vae_ckpt_path = vae_dir / "diffusion_pytorch_model.safetensors"
+    vae_config_path = vae_dir / "config.json"
+    with open(vae_config_path, 'r') as f:
+        vae_config = json.load(f)
+    vae = CausalVideoAutoencoder.from_config(vae_config)
+    vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
+    vae.load_state_dict(vae_state_dict)
+    return vae.cuda().to(torch.bfloat16)
 
-model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
-vae_local_path = Path("/opt/models/checkpoints/vae_training/causal_vvae_32x32x8_420m_cont_32/step_2296000")
-dtype = torch.float32
-vae = CausalVideoAutoencoder.from_pretrained(
-            pretrained_model_name_or_path=vae_local_path,
-            revision=False,
-            torch_dtype=torch.bfloat16,
-            load_in_8bit=False,
-).cuda()
-transformer_config_path = Path("/opt/txt2img/txt2img/config/transformer3d/xora_v1.2-L.json")
-transformer_config = Transformer3DModel.load_config(transformer_config_path)
-transformer = Transformer3DModel.from_config(transformer_config)
-transformer_local_path = Path("/opt/models/logs/v1.2-vae-mf-medHR-mr-cvae-nl/ckpt/01760000/model.pt")
-transformer_ckpt_state_dict = torch.load(transformer_local_path)
-transformer.load_state_dict(transformer_ckpt_state_dict, True)
-transformer = transformer.cuda()
-unet = transformer
-scheduler_config_path = Path("/opt/txt2img/txt2img/config/scheduler/RF_SD3_shifted.json")
-scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
-scheduler = RectifiedFlowScheduler.from_config(scheduler_config)
-patchifier = SymmetricPatchifier(patch_size=1)
-# text_encoder = T5EncoderModel.from_pretrained("t5-v1_1-xxl")
+def load_unet(unet_dir):
+    unet_ckpt_path = unet_dir / "diffusion_pytorch_model.safetensors"
+    unet_config_path = unet_dir / "config.json"
+    transformer_config = Transformer3DModel.load_config(unet_config_path)
+    transformer = Transformer3DModel.from_config(transformer_config)
+    unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
+    transformer.load_state_dict(unet_state_dict, strict=True)
+    return transformer.cuda()
 
-submodel_dict = {
-    "unet": unet,
-    "transformer": transformer,
-    "patchifier": patchifier,
-    "text_encoder": None,
-    "scheduler": scheduler,
-    "vae": vae,
+def load_scheduler(scheduler_dir):
+    scheduler_config_path = scheduler_dir / "scheduler_config.json"
+    scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
+    return RectifiedFlowScheduler.from_config(scheduler_config)
 
-}
+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')
+    args = parser.parse_args()
 
-pipeline = VideoPixArtAlphaPipeline.from_pretrained(model_name_or_path,
-                                                    safety_checker=None,
-            revision=None,
-            torch_dtype=dtype,
-            **submodel_dict,
-        )
+    # 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'
 
-num_inference_steps=20
-num_images_per_prompt=2
-guidance_scale=3
-height=512
-width=768
-num_frames=57
-frame_rate=25
-# sample = {
-#     "prompt": "A cat", # (B, L, E)
-#     'prompt_attention_mask': None, # (B , L)
-#     'negative_prompt': "Ugly deformed",
-#     'negative_prompt_attention_mask': None # (B , L)
-# }
+    # Load models
+    vae = load_vae(vae_dir)
+    unet = load_unet(unet_dir)
+    scheduler = load_scheduler(scheduler_dir)
 
-sample = torch.load("/opt/sample.pt")
-for _, item in sample.items():
-    if item is not None:
-        item = item.cuda()
+    # 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")
 
+    # Use submodels for the pipeline
+    submodel_dict = {
+        "transformer": unet,  # using unet for transformer
+        "patchifier": patchifier,
+        "scheduler": scheduler,
+        "text_encoder": text_encoder,
+        "tokenizer": tokenizer,
+        "vae": vae,
+    }
 
-images = pipeline(
-    num_inference_steps=num_inference_steps,
-    num_images_per_prompt=num_images_per_prompt,
-    guidance_scale=guidance_scale,
-    generator=None,
-    output_type="pt",
-    callback_on_step_end=None,
-    height=height,
-    width=width,
-    num_frames=num_frames,
-    frame_rate=frame_rate,
-    **sample,
-    is_video=True,
-    vae_per_channel_normalize=True,
-).images
+    pipeline = VideoPixArtAlphaPipeline(**submodel_dict).to("cuda")
 
-print()
+    # Sample input
+    num_inference_steps = 20
+    num_images_per_prompt = 2
+    guidance_scale = 3
+    height = 512
+    width = 768
+    num_frames = 57
+    frame_rate = 25
+    sample = {
+        "prompt": "A middle-aged man with glasses and a salt-and-pepper beard is driving a car and talking, gesturing with his right hand. "
+                  "The man is wearing a dark blue zip-up jacket and a light blue collared shirt. He is sitting in the driver's seat of a car with a black interior. The car is moving on a road with trees and bushes on either side. The man has a serious expression on his face and is looking straight ahead.",
+        'prompt_attention_mask': None,  # Adjust attention masks as needed
+        'negative_prompt': "Ugly deformed",
+        'negative_prompt_attention_mask': None
+    }
+
+    # 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,
+        output_type="pt",
+        callback_on_step_end=None,
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        frame_rate=frame_rate,
+        **sample,
+        is_video=True,
+        vae_per_channel_normalize=True,
+    ).images
+
+    print("Generated images (video frames).")
+
+if __name__ == "__main__":
+    main()

xora/models/autoencoders/causal_video_autoencoder.py

@@ -126,6 +126,13 @@ class CausalVideoAutoencoder(AutoencoderKLWrapper):
         return json.dumps(self.config.__dict__)
 
     def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
+        per_channel_statistics_prefix = "per_channel_statistics."
+        ckpt_state_dict = {
+            key: value
+            for key, value in state_dict.items()
+            if not key.startswith(per_channel_statistics_prefix)
+        }
+
         model_keys = set(name for name, _ in self.named_parameters())
 
         key_mapping = {
@@ -133,9 +140,8 @@ class CausalVideoAutoencoder(AutoencoderKLWrapper):
             "downsamplers.0": "downsample",
             "upsamplers.0": "upsample",
         }
-
         converted_state_dict = {}
-        for key, value in state_dict.items():
+        for key, value in ckpt_state_dict.items():
             for k, v in key_mapping.items():
                 key = key.replace(k, v)
 
@@ -147,6 +153,20 @@ class CausalVideoAutoencoder(AutoencoderKLWrapper):
 
         super().load_state_dict(converted_state_dict, strict=strict)
 
+        data_dict = {
+            key.removeprefix(per_channel_statistics_prefix): value
+            for key, value in state_dict.items()
+            if key.startswith(per_channel_statistics_prefix)
+        }
+        if len(data_dict) > 0:
+            self.register_buffer("std_of_means", data_dict["std-of-means"])
+            self.register_buffer(
+                "mean_of_means",
+                data_dict.get(
+                    "mean-of-means", torch.zeros_like(data_dict["std-of-means"])
+                ),
+            )
+
     def last_layer(self):
         if hasattr(self.decoder, "conv_out"):
             if isinstance(self.decoder.conv_out, nn.Sequential):