(vae_trt, cfg_denoiser, diffusion, val_data, clip_model, future_len,
history_len, cfg)
| 36 | |
| 37 | |
| 38 | def warmup(vae_trt, cfg_denoiser, diffusion, val_data, clip_model, future_len, |
| 39 | history_len, cfg): |
| 40 | |
| 41 | history_motion = val_data.normalize(get_zero_feature().unsqueeze(0).expand( |
| 42 | 1, history_len, -1).to(cfg.device)) |
| 43 | abs_pose = get_zero_abs_pose((1, ), device=cfg.device) |
| 44 | |
| 45 | def get_text_embedding(text: str, clip_model, device: str) -> torch.Tensor: |
| 46 | """Encode text using CLIP model.""" |
| 47 | try: |
| 48 | with torch.no_grad(): |
| 49 | import clip |
| 50 | text_tokens = clip.tokenize([text]).to(device) |
| 51 | text_embedding = clip_model.encode_text(text_tokens) |
| 52 | # text_embedding = text_embedding / text_embedding.norm( |
| 53 | # dim=-1, keepdim=True) |
| 54 | return text_embedding.float() |
| 55 | except Exception as e: |
| 56 | print(f"Warning: Failed to encode text '{text}': {e}") |
| 57 | return torch.zeros(1, 512, device=device, dtype=torch.float32) |
| 58 | |
| 59 | text_embedding = get_text_embedding("stand", clip_model, cfg.device) |
| 60 | future_motion, motion_dict, abs_pose = generate_next_motion( |
| 61 | vae=vae_trt, |
| 62 | denoiser=cfg_denoiser, |
| 63 | diffusion=diffusion, |
| 64 | val_data=val_data, |
| 65 | text_embedding=text_embedding, |
| 66 | history_motion=history_motion, |
| 67 | abs_pose=abs_pose, |
| 68 | future_len=future_len, |
| 69 | # cfg=cfg, |
| 70 | use_full_sample=cfg.use_full_sample, |
| 71 | guidance_scale=cfg.guidance_scale, |
| 72 | ) |
| 73 | # 因为第一次和第二次, history_motion的stride内存布局不一样, 会重新触发编译 |
| 74 | |
| 75 | # print(f"First history_motion strides: {history_motion.stride()}") |
| 76 | # print( |
| 77 | # f"Second history_motion strides: {future_motion[:, -history_len:, :].stride()}" |
| 78 | # ) |
| 79 | history_motion = future_motion[:, -history_len:, :] |
| 80 | future_motion, motion_dict, abs_pose = generate_next_motion( |
| 81 | vae=vae_trt, |
| 82 | denoiser=cfg_denoiser, |
| 83 | diffusion=diffusion, |
| 84 | val_data=val_data, |
| 85 | text_embedding=text_embedding, |
| 86 | history_motion=history_motion, |
| 87 | abs_pose=abs_pose, |
| 88 | future_len=future_len, |
| 89 | # cfg=cfg, |
| 90 | use_full_sample=cfg.use_full_sample, |
| 91 | guidance_scale=cfg.guidance_scale, |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | mujoco_to_isaaclab_reindex = [ |
no test coverage detected