Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -92,64 +92,107 @@ RESOLUTION_MAP = {
     "cogvideox-2b": (480, 720),
 }
 
-
-
-def load_cogvideox_flovd_FVSM_controlnet_pipeline(controlnet_path, backbone_path, device, dtype):
-    controlnet_sd = torch.load(controlnet_path, map_location='cpu')['module']
-
+def load_cogvideox_flovd_OMSM_lora_pipeline(omsm_path, backbone_path, transformer_lora_config, device, dtype):
     tokenizer = AutoTokenizer.from_pretrained(backbone_path, subfolder="tokenizer")
-    text_encoder = T5EncoderModel.from_pretrained(backbone_path, subfolder="text_encoder", device_map="auto", torch_dtype=dtype)
-    transformer = CogVideoXTransformer3DModel.from_pretrained(backbone_path, subfolder="transformer", device_map="auto", torch_dtype=dtype)
-    vae = AutoencoderKLCogVideoX.from_pretrained(backbone_path, subfolder="vae", device_map="auto", torch_dtype=dtype)
+    text_encoder = T5EncoderModel.from_pretrained(backbone_path, subfolder="text_encoder")
+    transformer = CogVideoXTransformer3DModel.from_pretrained(
+        backbone_path, subfolder="transformer", torch_dtype=dtype, device_map="auto"
+    )
+    vae = AutoencoderKLCogVideoX.from_pretrained(
+        backbone_path, subfolder="vae", torch_dtype=dtype, device_map="auto"
+    )
     scheduler = CogVideoXDPMScheduler.from_pretrained(backbone_path, subfolder="scheduler")
 
-    controlnet = CogVideoXControlnet.from_pretrained(backbone_path, subfolder="transformer", num_layers=6, out_proj_dim_factor=64, out_proj_dim_zero_init=True, notextinflow=True)
-    controlnet.eval()
-    controlnet.load_state_dict(controlnet_sd)
+    # 1) Load Lora weight
+    transformer.add_adapter(transformer_lora_config)
 
-    pipe = FloVDCogVideoXControlnetImageToVideoPipeline(
+    lora_state_dict = FloVDOMSMCogVideoXImageToVideoPipeline.lora_state_dict(omsm_path)
+    transformer_state_dict = {
+        f'{k.replace("transformer.", "")}': v
+        for k, v in lora_state_dict.items()
+        if k.startswith("transformer.")
+    }
+    incompatible_keys = set_peft_model_state_dict(transformer, transformer_state_dict, adapter_name="default")
+    if incompatible_keys is not None:
+        # check only for unexpected keys
+        unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
+        if unexpected_keys:
+            logger.warning(
+                f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
+                f" {unexpected_keys}. "
+            )
+
+    # 2) Load Other weight
+    load_path = os.path.join(omsm_path, "selected_blocks.safetensors")
+    if os.path.exists(load_path):
+        tensor_dict = load_file(load_path)
+
+        block_state_dicts = {}
+        for k, v in tensor_dict.items():
+            block_name, param_name = k.split(".", 1)
+            if block_name not in block_state_dicts:
+                block_state_dicts[block_name] = {}
+            block_state_dicts[block_name][param_name] = v
+
+        for block_name, state_dict in block_state_dicts.items():
+            if hasattr(transformer, block_name):
+                getattr(transformer, block_name).load_state_dict(state_dict)
+            else:
+                raise ValueError(f"Transformer has no attribute '{block_name}'")
+
+    pipe = FloVDOMSMCogVideoXImageToVideoPipeline(
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         vae=vae,
         transformer=transformer,
-        controlnet=controlnet,
         scheduler=scheduler,
-    ).to(device, dtype)
+    )
+
+    pipe.vae.enable_slicing()
+    pipe.vae.enable_tiling()
 
     return pipe
 
 
-def load_cogvideox_flovd_OMSM_lora_pipeline(omsm_path, backbone_path, transformer_lora_config, device, dtype):
+def load_cogvideox_flovd_FVSM_controlnet_pipeline(controlnet_path, backbone_path, device, dtype):
+    controlnet_sd = torch.load(controlnet_path, map_location='cpu')['module']
+
     tokenizer = AutoTokenizer.from_pretrained(backbone_path, subfolder="tokenizer")
-    text_encoder = T5EncoderModel.from_pretrained(backbone_path, subfolder="text_encoder", device_map="auto", torch_dtype=dtype)
-    transformer = CogVideoXTransformer3DModel.from_pretrained(backbone_path, subfolder="transformer", device_map="auto", torch_dtype=dtype)
-    vae = AutoencoderKLCogVideoX.from_pretrained(backbone_path, subfolder="vae", device_map="auto", torch_dtype=dtype)
+    text_encoder = T5EncoderModel.from_pretrained(backbone_path, subfolder="text_encoder")
+    transformer = CustomCogVideoXTransformer3DModel.from_pretrained(
+        backbone_path, subfolder="transformer", torch_dtype=dtype, device_map="auto"
+    )
+    vae = AutoencoderKLCogVideoX.from_pretrained(
+        backbone_path, subfolder="vae", torch_dtype=dtype, device_map="auto"
+    )
     scheduler = CogVideoXDPMScheduler.from_pretrained(backbone_path, subfolder="scheduler")
 
-    transformer.add_adapter(transformer_lora_config)
-    lora_state_dict = FloVDOMSMCogVideoXImageToVideoPipeline.lora_state_dict(omsm_path)
-    transformer_state_dict = {
-        f'{k.replace("transformer.", "")}': v
-        for k, v in lora_state_dict.items()
-        if k.startswith("transformer.")
+    additional_kwargs = {
+        'num_layers': 6,
+        'out_proj_dim_factor': 64,
+        'out_proj_dim_zero_init': True,
+        'notextinflow': True,
     }
-    set_peft_model_state_dict(transformer, transformer_state_dict, adapter_name="default")
+    controlnet = CogVideoXControlnet.from_pretrained(backbone_path, subfolder="transformer", **additional_kwargs)
+    controlnet.eval()
 
-    other_block_path = os.path.join(omsm_path, "selected_blocks.safetensors")
-    if os.path.exists(other_block_path):
-        tensor_dict = load_file(other_block_path)
-        for block_name, param in tensor_dict.items():
-            module_name, param_name = block_name.split(".", 1)
-            if hasattr(transformer, module_name):
-                getattr(transformer, module_name).load_state_dict({param_name: param}, strict=False)
+    missing, unexpected = controlnet.load_state_dict(controlnet_sd)
 
-    pipe = FloVDOMSMCogVideoXImageToVideoPipeline(
+    if len(missing) != 0 or len(unexpected) != 0:
+        print(f"Missing keys : {missing}")
+        print(f"Unexpected keys : {unexpected}")
+
+    pipe = FloVDCogVideoXControlnetImageToVideoPipeline(
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         vae=vae,
         transformer=transformer,
+        controlnet=controlnet,
         scheduler=scheduler,
-    ).to(device, dtype)
+    )
+
+    pipe.vae.enable_slicing()
+    pipe.vae.enable_tiling()
 
     return pipe