diff --git a/src/unifolm_wma/modules/networks/wma_model.py b/src/unifolm_wma/modules/networks/wma_model.py index 96e9df7..0d94fe4 100644 --- a/src/unifolm_wma/modules/networks/wma_model.py +++ b/src/unifolm_wma/modules/networks/wma_model.py @@ -848,15 +848,16 @@ class WMAModel(nn.Module): if not self.base_model_gen_only: ba, _, _ = x_action.shape + ts_state = timesteps[:ba] if b > 1 else timesteps + # Run action_unet and state_unet in parallel via CUDA streams + s_stream = torch.cuda.Stream() + s_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s_stream): + s_y = self.state_unet(x_state, ts_state, hs_a, + context_action[:2], **kwargs) a_y = self.action_unet(x_action, timesteps[:ba], hs_a, context_action[:2], **kwargs) - # Predict state - if b > 1: - s_y = self.state_unet(x_state, timesteps[:ba], hs_a, - context_action[:2], **kwargs) - else: - s_y = self.state_unet(x_state, timesteps, hs_a, - context_action[:2], **kwargs) + torch.cuda.current_stream().wait_stream(s_stream) else: a_y = torch.zeros_like(x_action) s_y = torch.zeros_like(x_state) diff --git a/unitree_z1_dual_arm_cleanup_pencils/case1/output.log b/unitree_z1_dual_arm_cleanup_pencils/case1/output.log index c94ef48..63a92cd 100644 --- a/unitree_z1_dual_arm_cleanup_pencils/case1/output.log +++ b/unitree_z1_dual_arm_cleanup_pencils/case1/output.log @@ -1,14 +1,14 @@ /mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/lightning_fabric/__init__.py:29: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81. __import__("pkg_resources").declare_namespace(__name__) -2026-02-10 06:42:14.444321: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. -2026-02-10 06:42:14.447338: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used. -2026-02-10 06:42:14.478442: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered -2026-02-10 06:42:14.478474: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered -2026-02-10 06:42:14.480279: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered -2026-02-10 06:42:14.488343: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used. -2026-02-10 06:42:14.488598: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. +2026-02-10 10:36:44.797852: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. +2026-02-10 10:36:44.801300: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used. +2026-02-10 10:36:44.837891: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered +2026-02-10 10:36:44.837946: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered +2026-02-10 10:36:44.839880: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered +2026-02-10 10:36:44.849073: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used. +2026-02-10 10:36:44.849365: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. -2026-02-10 06:42:15.109100: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT +2026-02-10 10:36:45.644793: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT [rank: 0] Global seed set to 123 /mnt/ASC1637/miniconda3/envs/unifolm-wma-o/lib/python3.10/site-packages/kornia/feature/lightglue.py:44: FutureWarning: `torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead. @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) @@ -25,7 +25,7 @@ INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k). INFO:root:Loaded ViT-H-14 model config. DEBUG:urllib3.connectionpool:https://hf-mirror.com:443 "HEAD /laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin HTTP/1.1" 302 0 INFO:root:Loading pretrained ViT-H-14 weights (laion2b_s32b_b79k). -/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:199: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. +/mnt/ASC1637/unifolm-world-model-action/scripts/evaluation/world_model_interaction.py:198: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. state_dict = torch.load(ckpt, map_location="cpu") >>> model checkpoint loaded. >>> Load pre-trained model ... @@ -116,7 +116,7 @@ DEBUG:PIL.Image:Importing WmfImagePlugin DEBUG:PIL.Image:Importing XbmImagePlugin DEBUG:PIL.Image:Importing XpmImagePlugin DEBUG:PIL.Image:Importing XVThumbImagePlugin - 12%|█▎ | 1/8 [01:08<07:56, 68.03s/it] 25%|██▌ | 2/8 [02:12<06:35, 65.91s/it] 38%|███▊ | 3/8 [03:16<05:26, 65.22s/it] 50%|█████ | 4/8 [04:21<04:19, 64.86s/it] 62%|██████▎ | 5/8 [05:25<03:13, 64.67s/it] 75%|███████▌ | 6/8 [06:29<02:09, 64.55s/it] 88%|████████▊ | 7/8 [07:34<01:04, 64.44s/it] 100%|██████████| 8/8 [08:38<00:00, 64.44s/it] 100%|██████████| 8/8 [08:38<00:00, 64.81s/it] + 12%|█▎ | 1/8 [01:06<07:46, 66.62s/it] 25%|██▌ | 2/8 [02:07<06:17, 62.97s/it] 38%|███▊ | 3/8 [03:07<05:08, 61.80s/it] 50%|█████ | 4/8 [04:07<04:05, 61.30s/it] 62%|██████▎ | 5/8 [05:08<03:03, 61.02s/it] 75%|███████▌ | 6/8 [06:08<02:01, 60.84s/it] 88%|████████▊ | 7/8 [07:09<01:00, 60.68s/it] 100%|██████████| 8/8 [08:09<00:00, 60.66s/it] 100%|██████████| 8/8 [08:09<00:00, 61.25s/it] >>>>>>>>>>>>>>>>>>>>>>>> >>> Step 1: generating actions ... >>> Step 1: interacting with world model ... @@ -140,6 +140,6 @@ DEBUG:PIL.Image:Importing XVThumbImagePlugin >>> Step 7: interacting with world model ... >>>>>>>>>>>>>>>>>>>>>>>> -real 9m42.052s -user 8m0.668s -sys 1m2.045s +real 9m16.436s +user 11m27.120s +sys 1m1.877s diff --git a/unitree_z1_dual_arm_cleanup_pencils/case1/run_world_model_interaction.sh b/unitree_z1_dual_arm_cleanup_pencils/case1/run_world_model_interaction.sh index a626d19..698fae2 100644 --- a/unitree_z1_dual_arm_cleanup_pencils/case1/run_world_model_interaction.sh +++ b/unitree_z1_dual_arm_cleanup_pencils/case1/run_world_model_interaction.sh @@ -22,6 +22,5 @@ dataset="unitree_z1_dual_arm_cleanup_pencils" --guidance_rescale 0.7 \ --perframe_ae \ --vae_dtype bf16 \ - --fast_policy_no_decode \ - --sampler_type unipc + --fast_policy_no_decode } 2>&1 | tee "${res_dir}/output.log"