Ckpt conversion: script + usage examples updated.

scripts/to_safetensors.py

@@ -0,0 +1,110 @@
+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, t5_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('--t5_path', '-t', type=str, default='t5/PixArt-XL-2-1024-MS/')
+    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,32 +5,46 @@ 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
+import safetensors.torch
+import json
 
+# Paths for the separate mode directories
+separate_dir = Path("/opt/models/xora-img2video")
+unet_dir = separate_dir / 'unet'
+vae_dir = separate_dir / 'vae'
+scheduler_dir = separate_dir / 'scheduler'
 
-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,
+# Load VAE from separate mode
+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_state_dict = safetensors.torch.load_file(vae_ckpt_path)
+vae = CausalVideoAutoencoder.from_pretrained_conf(
+    config=vae_config,
+    state_dict=vae_state_dict,
+    torch_dtype=torch.bfloat16
 ).cuda()
-transformer_config_path = Path("/opt/txt2img/txt2img/config/transformer3d/xora_v1.2-L.json")
-transformer_config = Transformer3DModel.load_config(transformer_config_path)
+
+# Load UNet (Transformer) from separate mode
+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)
-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)
+unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
+transformer.load_state_dict(unet_state_dict, strict=True)
 transformer = transformer.cuda()
 unet = transformer
-scheduler_config_path = Path("/opt/txt2img/txt2img/config/scheduler/RF_SD3_shifted.json")
+
+# Load Scheduler from separate mode
+scheduler_config_path = scheduler_dir / "scheduler_config.json"
 scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
 scheduler = RectifiedFlowScheduler.from_config(scheduler_config)
+
+# Patchifier (remains the same)
 patchifier = SymmetricPatchifier(patch_size=1)
-# text_encoder = T5EncoderModel.from_pretrained("t5-v1_1-xxl")
 
+# Use submodels for the pipeline
 submodel_dict = {
     "unet": unet,
     "transformer": transformer,
@@ -38,36 +52,33 @@ submodel_dict = {
     "text_encoder": None,
     "scheduler": scheduler,
     "vae": vae,
-
 }
 
+model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
 pipeline = VideoPixArtAlphaPipeline.from_pretrained(model_name_or_path,
                                                     safety_checker=None,
-            revision=None,
-            torch_dtype=dtype,
-            **submodel_dict,
-        )
+                                                    revision=None,
+                                                    torch_dtype=torch.float32,  # dtype adjusted
+                                                    **submodel_dict,
+                                                    ).to("cuda")
 
-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)
-# }
+num_inference_steps = 20
+num_images_per_prompt = 2
+guidance_scale = 3
+height = 512
+width = 768
+num_frames = 57
+frame_rate = 25
 
+# Assuming sample is a dict loaded from a .pt file
 sample = torch.load("/opt/sample.pt")
-for _, item in sample.items():
+for key, item in sample.items():
     if item is not None:
-        item = item.cuda()
+        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,
@@ -84,4 +95,4 @@ images = pipeline(
     vae_per_channel_normalize=True,
 ).images
 
-print()
+print("Generated video frames.")

xora/examples/text_to_video.py

@@ -6,69 +6,78 @@ from xora.schedulers.rf import RectifiedFlowScheduler
 from xora.pipelines.pipeline_video_pixart_alpha import VideoPixArtAlphaPipeline
 from pathlib import Path
 from transformers import T5EncoderModel
+import safetensors.torch
+import json
 
+# Paths for the separate mode directories
+separate_dir = Path("/opt/models/xora-txt2video")
+unet_dir = separate_dir / 'unet'
+vae_dir = separate_dir / 'vae'
+scheduler_dir = separate_dir / 'scheduler'
 
-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,
+# Load VAE from separate mode
+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_state_dict = safetensors.torch.load_file(vae_ckpt_path)
+vae = CausalVideoAutoencoder.from_pretrained_conf(
+    config=vae_config,
+    state_dict=vae_state_dict,
+    torch_dtype=torch.bfloat16
 ).cuda()
-transformer_config_path = Path("/opt/txt2img/txt2img/config/transformer3d/xora_v1.2-L.json")
-transformer_config = Transformer3DModel.load_config(transformer_config_path)
+
+# Load UNet (Transformer) from separate mode
+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)
-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)
+unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
+transformer.load_state_dict(unet_state_dict, strict=True)
 transformer = transformer.cuda()
 unet = transformer
-scheduler_config_path = Path("/opt/txt2img/txt2img/config/scheduler/RF_SD3_shifted.json")
+
+# Load Scheduler from separate mode
+scheduler_config_path = scheduler_dir / "scheduler_config.json"
 scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path)
 scheduler = RectifiedFlowScheduler.from_config(scheduler_config)
+
+# Patchifier (remains the same)
 patchifier = SymmetricPatchifier(patch_size=1)
-# text_encoder = T5EncoderModel.from_pretrained("t5-v1_1-xxl")
 
+# Use submodels for the pipeline
 submodel_dict = {
     "unet": unet,
     "transformer": transformer,
     "patchifier": patchifier,
-    "text_encoder": None,
     "scheduler": scheduler,
     "vae": vae,
-
 }
-
+model_name_or_path = "PixArt-alpha/PixArt-XL-2-1024-MS"
 pipeline = VideoPixArtAlphaPipeline.from_pretrained(model_name_or_path,
                                                     safety_checker=None,
             revision=None,
-            torch_dtype=dtype,
+            torch_dtype=torch.float32,
             **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()
-
+        ).to("cuda")
 
+# 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,
@@ -85,4 +94,4 @@ images = pipeline(
     vae_per_channel_normalize=True,
 ).images
 
-print()
+print("Generated images (video frames).")

xora/models/autoencoders/causal_video_autoencoder.py

@@ -41,6 +41,35 @@ class CausalVideoAutoencoder(AutoencoderKLWrapper):
 
         return video_vae
 
+    @classmethod
+    def from_pretrained_conf(cls, config, state_dict, torch_dtype=torch.float32):
+        video_vae = cls.from_config(config)
+        video_vae.to(torch_dtype)
+
+        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)
+        }
+        video_vae.load_state_dict(ckpt_state_dict)
+
+        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:
+            video_vae.register_buffer("std_of_means", data_dict["std-of-means"])
+            video_vae.register_buffer(
+                "mean_of_means",
+                data_dict.get(
+                    "mean-of-means", torch.zeros_like(data_dict["std-of-means"])
+                ),
+            )
+
+        return video_vae
+
     @staticmethod
     def from_config(config):
         assert config["_class_name"] == "CausalVideoAutoencoder", "config must have _class_name=CausalVideoAutoencoder"