| 64 | |
| 65 | |
| 66 | def cache_video( |
| 67 | tensor, |
| 68 | save_file: str, |
| 69 | fps=30, |
| 70 | suffix=".mp4", |
| 71 | nrow=8, |
| 72 | normalize=True, |
| 73 | value_range=(-1, 1), |
| 74 | retry=5, |
| 75 | ): |
| 76 | save_dir = os.path.dirname(save_file) |
| 77 | try: |
| 78 | if not os.path.exists(save_dir): |
| 79 | os.makedirs(save_dir, exist_ok=True) |
| 80 | except Exception as e: |
| 81 | logger.error(f"Failed to create directory: {save_dir}, error: {e}") |
| 82 | return None |
| 83 | |
| 84 | cache_file = save_file |
| 85 | |
| 86 | # save to cache |
| 87 | error = None |
| 88 | for _ in range(retry): |
| 89 | try: |
| 90 | # preprocess |
| 91 | tensor = tensor.clamp(min(value_range), max(value_range)) # type: ignore |
| 92 | tensor = torch.stack( |
| 93 | [torchvision.utils.make_grid(u, nrow=nrow, normalize=normalize, value_range=value_range) for u in tensor.unbind(2)], |
| 94 | dim=1, |
| 95 | ).permute(1, 2, 3, 0) |
| 96 | tensor = (tensor * 255).type(torch.uint8).cpu() |
| 97 | |
| 98 | # write video |
| 99 | writer = imageio.get_writer(cache_file, fps=fps, codec="libx264", quality=8) |
| 100 | for frame in tensor.numpy(): |
| 101 | writer.append_data(frame) |
| 102 | writer.close() |
| 103 | del tensor |
| 104 | torch.cuda.empty_cache() |
| 105 | return cache_file |
| 106 | except Exception as e: |
| 107 | error = e |
| 108 | continue |
| 109 | else: |
| 110 | logger.info(f"cache_video failed, error: {error}", flush=True) |
| 111 | return None |
| 112 | |
| 113 | |
| 114 | def vae_to_comfyui_image(vae_output: torch.Tensor) -> torch.Tensor: |