[CI] Fix SSIM and transformers CI (#564)

Author: Edenzzzz

fastvideo/v1/models/loader/component_loader.py

@@ -252,7 +252,7 @@ def load_model(self,
                    fastvideo_args: FastVideoArgs,
                    dtype: str = "fp16"):
         use_cpu_offload = fastvideo_args.text_encoder_offload and len(
-            model_config._fsdp_shard_conditions) > 0
+            getattr(model_config, "_fsdp_shard_conditions", [])) > 0
 
         if fastvideo_args.text_encoder_offload:
             target_device = torch.device("cpu")

fastvideo/v1/models/loader/fsdp_load.py

@@ -106,8 +106,8 @@ def maybe_load_fsdp_model(
                 fsdp_shard_conditions=model._fsdp_shard_conditions,
                 pin_cpu_memory=pin_cpu_memory)
 
-    weight_iterator = safetensors_weights_iterator(weight_dir_list,
-                                                   async_broadcast=True)
+    weight_iterator = safetensors_weights_iterator(
+        weight_dir_list, to_cpu=cpu_offload, async_broadcast=not cpu_offload)
     param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping)
     load_model_from_full_model_state_dict(
         model,

fastvideo/v1/tests/transformers/test_hunyuanvideo.py

@@ -80,7 +80,7 @@ def test_hunyuanvideo_distributed():
 
     # Initialize with identical weights
     model = initialize_identical_weights(model, seed=42)
-    shard_model(model, cpu_offload=False,
+    shard_model(model, cpu_offload=True,
                 reshard_after_forward=True,
                 fsdp_shard_conditions=model._fsdp_shard_conditions
                 )
@@ -95,7 +95,7 @@ def test_hunyuanvideo_distributed():
 
     # Move to GPU based on local rank (0 or 1 for 2 GPUs)
     device = torch.device(f"cuda:0")
-    model = model.to(device)
+    model = model
 
     batch_size = 1
     seq_len = 3

fastvideo/v1/tests/transformers/test_hunyuanvideo_load.py

@@ -62,7 +62,7 @@ def test_hunyuanvideo_distributed():
 
     precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         use_cpu_offload=False,
+                         use_cpu_offload=True,
                          pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str))
     args.device = torch.device(f"cuda:{LOCAL_RANK}")
 

fastvideo/v1/tests/transformers/test_wanvideo.py

@@ -34,12 +34,12 @@ def test_wan_transformer():
     precision = torch.bfloat16
     precision_str = "bf16"
     args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         use_cpu_offload=False,
+                         use_cpu_offload=True,
                          pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
     args.device = device
 
     loader = TransformerLoader()
-    model2 = loader.load(TRANSFORMER_PATH, "", args).to(device, dtype=precision)
+    model2 = loader.load(TRANSFORMER_PATH, "", args).to(dtype=precision)
 
     model1 = WanTransformer3DModel.from_pretrained(
         TRANSFORMER_PATH, device=device,