Update demo.py

demo.py

@@ -253,44 +253,111 @@ class TrajCrafter:
         pose_t = poses
         return pose_s, pose_t, K
 
-    def setup_diffusion(self,opts):
-        # transformer = CrossTransformer3DModel.from_pretrained_cus(opts.transformer_path).to(opts.weight_dtype)
-        transformer = CrossTransformer3DModel.from_pretrained(opts.transformer_path).to(opts.weight_dtype)
-        # transformer = transformer.to(opts.weight_dtype)
+    # def setup_diffusion(self,opts):
+    #     # transformer = CrossTransformer3DModel.from_pretrained_cus(opts.transformer_path).to(opts.weight_dtype)
+    #     transformer = CrossTransformer3DModel.from_pretrained(opts.transformer_path).to(opts.weight_dtype)
+    #     # transformer = transformer.to(opts.weight_dtype)
+    #     vae = AutoencoderKLCogVideoX.from_pretrained(
+    #         opts.model_name, 
+    #         subfolder="vae"
+    #     ).to(opts.weight_dtype)
+    #     text_encoder = T5EncoderModel.from_pretrained(
+    #         opts.model_name, subfolder="text_encoder", torch_dtype=opts.weight_dtype
+    #     )
+    #     # Get Scheduler
+    #     Choosen_Scheduler  = {
+    #         "Euler": EulerDiscreteScheduler,
+    #         "Euler A": EulerAncestralDiscreteScheduler,
+    #         "DPM++": DPMSolverMultistepScheduler, 
+    #         "PNDM": PNDMScheduler,
+    #         "DDIM_Cog": CogVideoXDDIMScheduler,
+    #         "DDIM_Origin": DDIMScheduler,
+    #     }[opts.sampler_name]
+    #     scheduler = Choosen_Scheduler.from_pretrained(
+    #         opts.model_name, 
+    #         subfolder="scheduler"
+    #     )
+
+    #     self.pipeline = TrajCrafter_Pipeline.from_pretrained(
+    #         opts.model_name,
+    #         vae=vae,
+    #         text_encoder=text_encoder,
+    #         transformer=transformer,
+    #         scheduler=scheduler,
+    #         torch_dtype=opts.weight_dtype
+    #     )
+
+    #     if opts.low_gpu_memory_mode:
+    #         self.pipeline.enable_sequential_cpu_offload()
+    #     else:
+    #         self.pipeline.enable_model_cpu_offload()
+
+    def setup_diffusion(self, opts):
+        import torch
+    
+        # 1) 选择设备
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+        # （可选）在 CPU 跑时避免 efficient attention 报错；在 CUDA 上也无害
+        try:
+            torch.backends.cuda.enable_flash_sdp(False)
+            torch.backends.cuda.enable_mem_efficient_sdp(True)
+            torch.backends.cuda.enable_math_sdp(True)
+        except Exception:
+            pass
+    
+        # 2) 加载/放置子模块到 device + dtype
+        # 注意：原代码只 .to(dtype)，未指定 device；这里补齐
+        transformer = CrossTransformer3DModel.from_pretrained(opts.transformer_path)
+        transformer = transformer.to(device=device, dtype=opts.weight_dtype)
+    
         vae = AutoencoderKLCogVideoX.from_pretrained(
-            opts.model_name, 
-            subfolder="vae"
-        ).to(opts.weight_dtype)
+            opts.model_name,
+            subfolder="vae",
+            # 仅指定 dtype；后面统一 .to(device)
+            # 某些 from_pretrained 不支持 device 形参
+        ).to(dtype=opts.weight_dtype).to(device)
+    
         text_encoder = T5EncoderModel.from_pretrained(
-            opts.model_name, subfolder="text_encoder", torch_dtype=opts.weight_dtype
-        )
-        # Get Scheduler
+            opts.model_name,
+            subfolder="text_encoder",
+            torch_dtype=opts.weight_dtype,
+        ).to(device)
+    
+        # 3) 调度器照旧
         Choosen_Scheduler  = {
             "Euler": EulerDiscreteScheduler,
             "Euler A": EulerAncestralDiscreteScheduler,
-            "DPM++": DPMSolverMultistepScheduler, 
+            "DPM++": DPMSolverMultistepScheduler,
             "PNDM": PNDMScheduler,
             "DDIM_Cog": CogVideoXDDIMScheduler,
             "DDIM_Origin": DDIMScheduler,
         }[opts.sampler_name]
         scheduler = Choosen_Scheduler.from_pretrained(
-            opts.model_name, 
+            opts.model_name,
             subfolder="scheduler"
         )
-
+    
+        # 4) 组装 pipeline，并确保在正确 device/dtype
         self.pipeline = TrajCrafter_Pipeline.from_pretrained(
             opts.model_name,
             vae=vae,
             text_encoder=text_encoder,
             transformer=transformer,
             scheduler=scheduler,
-            torch_dtype=opts.weight_dtype
+            torch_dtype=opts.weight_dtype,
         )
-
+    
+        # Offload 策略：
+        # - 如果你机器有足够显存，建议直接 to(device) 获得最稳的行为
+        # - 如果显存紧张，再启用 offload（需要 accelerate 支持）
         if opts.low_gpu_memory_mode:
-            self.pipeline.enable_sequential_cpu_offload()
-        else:
+            # 这两种 offload 会在计算时把块迁移到 GPU，空闲时回收；加速略低但更省显存
+            # 二选一：根据你之前的使用习惯保留其一
+            # self.pipeline.enable_sequential_cpu_offload()
             self.pipeline.enable_model_cpu_offload()
+        else:
+            self.pipeline.to(device)
 
     def run_gradio(self,input_video, stride, radius_scale, pose, steps, seed):
         frames = read_video_frames(input_video, self.opts.video_length, stride,self.opts.max_res)