Update inference/flovd_demo.py

inference/flovd_demo.py

@@ -94,104 +94,63 @@ RESOLUTION_MAP = {
 
 
 
-
 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")
-    transformer = CustomCogVideoXTransformer3DModel.from_pretrained(backbone_path, subfolder="transformer")
-    vae = AutoencoderKLCogVideoX.from_pretrained(backbone_path, subfolder="vae")
+    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)
     scheduler = CogVideoXDPMScheduler.from_pretrained(backbone_path, subfolder="scheduler")
-    
-    additional_kwargs = {
-        'num_layers': 6,
-        'out_proj_dim_factor': 64,
-        'out_proj_dim_zero_init': True,
-        'notextinflow': True,
-    }
-    controlnet = CogVideoXControlnet.from_pretrained(backbone_path, subfolder="transformer", **additional_kwargs)
+
+    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()
-    
-    missing, unexpected = controlnet.load_state_dict(controlnet_sd)
-    
-    if len(missing) != 0 or len(unexpected) != 0:
-        print(f"Missing keys : {missing}")
-        print(f"Unexpected keys : {unexpected}")
-        
+    controlnet.load_state_dict(controlnet_sd)
+
     pipe = FloVDCogVideoXControlnetImageToVideoPipeline(
-            tokenizer=tokenizer,
-            text_encoder=text_encoder,
-            vae=vae,
-            transformer=transformer,
-            controlnet=controlnet,
-            scheduler=scheduler,
-    )
-    
-    # pipe.enable_model_cpu_offload(device=device)
-    pipe = pipe.to(device, dtype)
-    
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        vae=vae,
+        transformer=transformer,
+        controlnet=controlnet,
+        scheduler=scheduler,
+    ).to(device, dtype)
+
     return pipe
 
+
 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")
-    transformer = CogVideoXTransformer3DModel.from_pretrained(backbone_path, subfolder="transformer")
-    vae = AutoencoderKLCogVideoX.from_pretrained(backbone_path, subfolder="vae")
+    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)
     scheduler = CogVideoXDPMScheduler.from_pretrained(backbone_path, subfolder="scheduler")
 
-    # 1) Load Lora weight
     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.")
     }
-    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}. "
-            )
+    set_peft_model_state_dict(transformer, transformer_state_dict, adapter_name="default")
+
+    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)
 
-    # 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,
-            scheduler=scheduler,
-    )
-    
-    # pipe.load_lora_weights(omsm_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1")
-    # pipe.fuse_lora(components=["transformer"], lora_scale=1.0)
-    
-    # pipe.enable_model_cpu_offload(device=device)
-    pipe = pipe.to(device, dtype)
-    
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        vae=vae,
+        transformer=transformer,
+        scheduler=scheduler,
+    ).to(device, dtype)
+
     return pipe