Coverage for wrapper/hunyuanframepackvae/wrapper_hunyuanframepackvae.py: 63%
170 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Handle video VAE encoding and decoding using the Hunyuan Framepack VAE model.
3"""
4import logging
5import os
6import tempfile
7import aiofiles
8import asyncio
10import torch
12from typing import override
13from typing import Union
14from typing import Optional
15from typing import Dict
16from typing import Any
18from torch import inference_mode
20from wrapper_model import ModelGeneration
22from model_timing import GenTimer
24from diffusers import AutoencoderKLHunyuanVideo
26from media_utils import base64_to_tensor
27from media_utils import save_bcthw_as_mp4
30class HunyuanFramepackVAEGeneration(ModelGeneration):
31 """Handle video VAE encoding and decoding using the Hunyuan Framepack VAE model."""
33 def __init__(
34 self,
35 param_dtype: torch.dtype = torch.float16,
36 enable_tiling: bool = False,
37 enable_slicing: bool = False,
38 ) -> None:
39 super().__init__("hunyuanframepackvae")
41 self.param_dtype = param_dtype
42 self.enable_tiling = enable_tiling
43 self.enable_slicing = enable_slicing
45 # Parallelism
46 self.GPU = None
47 if torch.cuda.is_available():
48 self.GPU = torch.cuda.get_device_name(0)
50 # Model features
51 self.latent_channels = 16 # Latent channels for Hunyuan VAE
52 self.vae_stride = (4, 8, 8) # time, height, width
53 self.FPS = 30 # This is technically a constant for the model
55 # Model components
56 self.vae: Optional[AutoencoderKLHunyuanVideo] = None
58 def __del__(self) -> None:
59 # Clean models
60 if self.vae is not None:
61 del self.vae
62 super().__del__()
64 def init_parallelism(self) -> None:
65 self.load_timer.start("torch_dist")
66 self.rank = int(os.getenv("RANK", 0))
67 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
68 self.world_size = int(os.getenv("WORLD_SIZE", 1))
69 if self.world_size > 1:
70 logging.warning(f"[{self.rank}] No distributed mode available.")
72 if torch.cuda.is_available():
73 self.device_id = self.local_rank
74 self.device = torch.device(f"cuda:{self.device_id}")
75 torch.cuda.set_device(self.local_rank)
76 else:
77 # Running on CPU is very slow, but it is supported
78 self.device_id = 0
79 self.device = torch.device("cpu")
80 # MAX_NUM_CPUS = 16
81 # num_threads = min(MAX_NUM_CPUS, os.cpu_count() or 1)
82 # torch.set_num_threads(num_threads)
83 num_threads = torch.get_num_threads()
84 logging.warning(f"CUDA is not available. Running VAE with {num_threads} CPU threads.")
86 self.load_timer.end("torch_dist")
88 def load_model(self) -> None:
89 prev_memory = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0
90 self.load_timer.start("vae")
91 self.vae = AutoencoderKLHunyuanVideo.from_pretrained(
92 "hunyuanvideo-community/HunyuanVideo",
93 subfolder="vae",
94 torch_dtype=self.param_dtype,
95 ).to(self.device)
96 assert self.vae is not None
97 self.vae.eval().requires_grad_(False) # type: ignore[union-attr]
99 if not self.enable_tiling:
100 logging.info(f"[{self.rank}] Disabling tiling for VAE.")
101 self.vae.disable_tiling() # type: ignore[union-attr]
102 else:
103 logging.info(f"[{self.rank}] Enabling tiling for VAE.")
104 self.vae.enable_tiling() # type: ignore[union-attr]
106 if not self.enable_slicing:
107 logging.info(f"[{self.rank}] Disabling slicing for VAE.")
108 self.vae.disable_slicing() # type: ignore[union-attr]
109 else:
110 logging.info(f"[{self.rank}] Enabling slicing for VAE.")
111 self.vae.enable_slicing() # type: ignore[union-attr]
112 self.load_timer.end("vae")
114 if torch.cuda.is_available():
115 diff_memory = torch.cuda.memory_allocated() - prev_memory
116 logging.info(f"[{self.rank}] VAE memory allocated: {diff_memory / 1024 / 1024 ** 2:.2f} GB.")
118 def init_model_parallelism(self) -> None:
119 if self.world_size > 1:
120 logging.warning(f"[{self.rank}] No distributed mode available for Hunyuan VAE.")
122 def model_compile(self) -> None:
123 if not self.torch_compile:
124 return
126 logging.info(f"[{self.rank}] Compiling VAE with torch.compile().")
127 self.load_timer.start("vae_compile")
128 assert self.vae is not None
129 self.vae = torch.compile( # type: ignore[call-overload]
130 self.vae,
131 mode="max-autotune-no-cudagraphs",
132 )
133 self.load_timer.end("vae_compile")
135 def _assert_model_init(self) -> None:
136 super()._assert_model_init()
137 assert self.vae is not None
139 def _assert_args(
140 self,
141 latents: torch.Tensor,
142 ) -> None:
143 if latents is None:
144 raise ValueError("Latents cannot be None.")
145 if not isinstance(latents, torch.Tensor):
146 raise TypeError(f"Expected latents to be a torch.Tensor, got {type(latents)}.")
147 if latents.ndim != 5:
148 raise ValueError(f"Expected latents with 5D [B, C, T, H, W], got {latents.ndim} dimensions.")
149 if latents.shape[1] != 16:
150 raise ValueError(f"Expected latents with 16 channels, got {latents.shape[1]} channels.")
151 if latents.shape[2] <= 0:
152 raise ValueError(f"Latents must have a positive number of frames, got {latents.shape[2]} frames.")
154 @inference_mode()
155 async def warmup(self) -> None:
156 logging.info(f"[{self.rank}] Warmup for Hunyuan Framepack VAE generation.")
157 latents = torch.randn(
158 (1, 16, 4, 64, 64), # B, C, T, H, W
159 device=self.device,
160 dtype=self.param_dtype
161 )
162 await self.generate(latents)
164 @override
165 @inference_mode()
166 async def generate(
167 self,
168 latents: torch.Tensor,
169 job_id: Optional[str] = None,
170 output_type: str = "tensor", # "tensor", "video_binary", "video_path"
171 ) -> Union[torch.Tensor, str, bytes, None]:
172 return await self.vae_decode(
173 latents,
174 job_id=job_id,
175 output_type=output_type)
177 @inference_mode()
178 async def vae_decode(
179 self,
180 latents: torch.Tensor,
181 job_id: Optional[str] = None,
182 output_type: str = "tensor", # "tensor", "video_binary", "video_path"
183 ) -> Union[torch.Tensor, str, bytes, None]:
184 """
185 Latent -> Pixels.
186 """
187 gen_timer = self._new_gen_timer(job_id)
189 self._assert_model_init()
190 assert self.vae is not None
191 self._assert_args(latents)
193 self.running = True
195 try:
196 gen_timer.start("vae_decoder")
197 latents = latents / self.vae.config.scaling_factor # type: ignore[attr-defined]
198 latents = latents.to(self.vae.device, dtype=self.param_dtype) # type: ignore[attr-defined]
199 pixels = await asyncio.to_thread(
200 self.vae.decode, # type: ignore[attr-defined]
201 latents)
202 pixels = pixels.sample
203 gen_timer.end("vae_decoder")
205 return await self._output_video(
206 job_id,
207 gen_timer,
208 pixels,
209 output_type)
210 finally:
211 self.running = False
212 gen_timer.end("total")
214 async def _output_video(
215 self,
216 job_id: Optional[str],
217 gen_timer: GenTimer,
218 pixels: torch.Tensor,
219 output_type: str = "tensor", # "tensor", "video_binary", "video_path"
220 ) -> Optional[Union[torch.Tensor, str, bytes]]:
221 # TODO use the one in HunyuanFramePackBase
222 gen_timer.start("output")
223 try:
224 if output_type == "tensor":
225 return pixels
227 if output_type in ("video_binary", "video_path"):
228 if not job_id:
229 video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
230 else:
231 video_path = f"/tmp/{job_id}.mp4"
232 video_path = save_bcthw_as_mp4(
233 pixels,
234 video_path,
235 fps=self.FPS)
236 if output_type == "video_path":
237 return video_path
239 # video_binary
240 async with aiofiles.open(video_path, "rb") as f:
241 video_binary = await f.read()
242 return video_binary
244 logging.error(f"Unknown output type: {output_type}")
245 return None
246 finally:
247 gen_timer.end("output")
249 @inference_mode()
250 def vae_encode(
251 self,
252 pixels: torch.Tensor,
253 job_id: Optional[str] = None,
254 ) -> torch.Tensor:
255 """
256 Pixels -> Latent.
257 """
258 gen_timer = self._new_gen_timer(job_id)
260 self._assert_model_init()
261 assert self.vae is not None
262 assert pixels is not None
263 assert isinstance(pixels, torch.Tensor)
264 assert pixels.ndim == 5 # B, C, T, H, W
265 assert pixels.shape[1] == 3 # RGB channels
267 self.running = True
269 try:
270 gen_timer.start("vae_encoder")
271 pixels = pixels.to(self.device, dtype=self.param_dtype)
272 latents = self.vae.encode(pixels).latent_dist.sample() # type: ignore[attr-defined]
273 latents = latents * self.vae.config.scaling_factor # type: ignore[attr-defined]
274 gen_timer.end("vae_encoder")
275 return latents
276 finally:
277 self.running = False
278 gen_timer.end("total")
280 async def get_rest_args(self, data_json: Dict[str, str]) -> Dict[str, Any]:
281 if data_json is None or not isinstance(data_json, dict):
282 raise ValueError("Missing JSON body")
284 latents_base64 = data_json.get("latents", None)
285 if latents_base64 is None:
286 raise ValueError("Missing 'latents' parameter")
287 latents = base64_to_tensor(latents_base64)
289 return {
290 "task": self.model_name,
291 "args": {
292 "latents": latents
293 }
294 }
296 def get_health(self) -> Dict[str, Any]:
297 ret = super().get_health()
298 ret.update({
299 "gpu": self.GPU,
300 "rank": self.rank,
301 "local_rank": self.local_rank,
302 "world_size": self.world_size,
303 "torch_compile": self.torch_compile,
304 "dtype": str(self.param_dtype),
305 "vae_stride": self.vae_stride,
306 "latent_channels": self.latent_channels,
307 "fps": self.FPS,
308 })
309 return ret