Refactor cpu_offload checks to ensure attribute existence before accessing (#115)

lightx2v/models/input_encoders/hf/xlm_roberta/model.py

@@ -434,7 +434,7 @@ class CLIPModel:
         logger.info(f"End Loading weights from {self.checkpoint_path}")
 
     def visual(self, videos, args):
-        if args.cpu_offload:
+        if hasattr(args, "cpu_offload") and args.cpu_offload:
             self.to_cuda()
         # preprocess
         size = (self.model.image_size,) * 2
@@ -445,7 +445,7 @@ class CLIPModel:
         with torch.amp.autocast("cuda", dtype=self.dtype):
             out = self.model.visual(videos, use_31_block=True)
 
-        if args.cpu_offload:
+        if hasattr(args, "cpu_offload") and args.cpu_offload:
             self.to_cpu()
         return out
 

lightx2v/models/video_encoders/hf/wan/vae.py

@@ -868,7 +868,7 @@ class WanVAE:
         """
         videos: A list of videos each with shape [C, T, H, W].
         """
-        if args.cpu_offload:
+        if hasattr(args, "cpu_offload") and args.cpu_offload:
             self.to_cuda()
 
         if self.use_tiling:
@@ -876,7 +876,7 @@ class WanVAE:
         else:
             out = [self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) for u in videos]
 
-        if args.cpu_offload:
+        if hasattr(args, "cpu_offload") and args.cpu_offload:
             self.to_cpu()
         return out