Coverage for wrapper/hunyuanimage/wrapper_hunyuanimage.py: 84%
125 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Wrapper for HunyuanImage model.
3"""
4import logging
5import os
6import asyncio
7import datetime
9from typing import override
10from typing import Any
11from typing import Dict
12from typing import Optional
14from PIL import Image
16import torch
17import torch.distributed as dist
18from torch import inference_mode
20from wrapper_model import ModelGeneration
22from transformers import AutoModelForCausalLM
23# from hunyuan_image_3_pipeline import HunyuanImage3Text2ImagePipeline
24from diffusers.pipelines.pipeline_utils import DiffusionPipeline
26from xfuser.config import EngineConfig
29class HunyuanImageGeneration(ModelGeneration):
30 """Handle image generation using the HunyuanImage model."""
32 # HF_MODEL_NAME = "tencent/HunyuanImage-3.0"
33 HF_MODEL_NAME = "./HunyuanImage-3"
34 MAX_LOG_TEXT_LEN = 64
36 def __init__(
37 self,
38 model_name: str = "hunyuanimage",
39 engine_config: EngineConfig = None,
40 param_dtype: torch.dtype = torch.bfloat16,
41 ) -> None:
42 super().__init__(model_name)
44 self.engine_config = engine_config
45 if self.engine_config is not None:
46 self.torch_compile = self.engine_config.runtime_config.use_torch_compile
47 self.param_dtype = param_dtype
49 # Model components
50 self.model: Optional[AutoModelForCausalLM] = None
51 self.pipeline: Optional[DiffusionPipeline] = None
53 def __del__(self) -> None:
54 if self.model is not None:
55 self.model = None
57 def init_parallelism(self) -> None:
58 self.load_timer.start("torch_dist")
60 self.rank = int(os.getenv("RANK", 0))
61 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
62 self.world_size = int(os.getenv("WORLD_SIZE", 1))
64 self.device_id = self.local_rank
66 if not torch.cuda.is_available():
67 self.device_id = 0
68 self.device = torch.device("cpu")
69 logging.warning("CUDA is not available. Running on CPU.")
70 self.load_timer.end("torch_dist")
71 return # Single GPU mode, no parallelism needed
73 self.device = torch.device(f"cuda:{self.device_id}")
75 torch.cuda.set_device(self.local_rank)
77 if self.world_size <= 1:
78 self.load_timer.end("torch_dist")
79 return # Single GPU mode, no parallelism needed
81 if not dist.is_initialized():
82 dist.init_process_group(
83 backend="nccl",
84 init_method="env://",
85 rank=self.rank,
86 world_size=self.world_size,
87 timeout=datetime.timedelta(hours=24), # Prevent NCCL timeout
88 )
90 self.load_timer.end("torch_dist")
92 if not dist.is_initialized():
93 raise RuntimeError("Distributed process group not initialized")
95 def init_model_parallelism(self) -> None:
96 logging.info(f"[{self.rank}] Hunyuan Image parallelism.")
98 def load_model(self) -> None:
99 assert torch.cuda.is_available()
101 if self.rank > 0:
102 logging.info(f"[{self.rank}] Model loaded only on rank 0.")
103 return
105 self.load_timer.start("model")
106 self.model = AutoModelForCausalLM.from_pretrained( # type: ignore[assignment]
107 self.HF_MODEL_NAME,
108 attn_implementation="sdpa", # Use "flash_attention_2" if FlashAttention is installed
109 trust_remote_code=True,
110 torch_dtype=self.param_dtype,
111 device_map="auto",
112 moe_impl="eager", # Use "flashinfer" if FlashInfer is installed
113 # low_cpu_mem_usage=True, # TODO ?
114 moe_drop_tokens=True,
115 ) # nosec B615 - local path
116 assert self.model is not None
117 self.model.load_tokenizer(self.HF_MODEL_NAME) # type: ignore[attr-defined]
118 self.pipeline = self.model.pipeline # type: ignore[attr-defined]
119 self.load_timer.end("model")
121 def model_compile(self) -> None:
122 """Compile the model using torch.compile if enabled."""
123 if not self.torch_compile:
124 return
126 if self.model:
127 self.load_timer.start("dit_compile")
128 self.model = torch.compile( # type: ignore[call-overload]
129 self.model,
130 mode="max-autotune-no-cudagraphs"
131 )
132 self.load_timer.end("dit_compile")
134 def _assert_args(
135 self,
136 height: int,
137 width: int,
138 ) -> None:
139 # height_latent = height // self.pipeline.vae_scale_factor
140 # width_latent = width // self.pipeline.vae_scale_factor
141 # self.model.vae is AutoencoderKLConv3D
142 assert self.model is not None
143 vae_config = self.model.vae.config # type: ignore[attr-defined]
144 if width % vae_config.ffactor_spatial != 0:
145 raise ValueError(f"Width {width} not supported. Must be multiple of {vae_config.ffactor_spatial}.")
146 if height % vae_config.ffactor_spatial != 0:
147 raise ValueError(f"Height {height} not supported. Must be multiple of {vae_config.ffactor_spatial}.")
148 """
149 height x width:
150 ("1:1", "1024x1024"),
151 ("4:3", "896x1152"),
152 ("3:4", "1152x896"),
153 ("16:9", "768x1280"),
154 ("9:16", "1280x768"),
155 ("21:9", "640x1408"),
156 """
157 if width * height > 1024 * 1024:
158 raise ValueError(f"{width}x{height} too large. Max is 1024 x 1024.")
160 def _assert_model_init(self) -> None:
161 super()._assert_model_init()
162 if self.model is None:
163 raise ValueError("HunyuanImage model not loaded.")
165 @inference_mode()
166 async def warmup(self) -> None:
167 logging.info(f"[{self.rank}] Warmup for Hunyuan Image generation.")
168 await self.generate(
169 height=1024, # 1:1
170 width=1024,
171 prompt="A warmup image to initialize the model.",
172 sampling_steps=5, # It needs at least 5 steps to warm up properly
173 )
175 @override
176 @inference_mode()
177 async def generate(
178 self,
179 height: int,
180 width: int,
181 prompt: str,
182 sampling_steps: int = 25,
183 cfg: float = 0.5,
184 seed: Optional[int] = None,
185 job_id: Optional[str] = None,
186 ) -> Optional[Image.Image]:
187 """Generate an image using HunyuanImage 3."""
188 if self.rank > 0:
189 logging.info(f"[{self.rank}] Image generation only rank 0.")
190 return None
192 gen_timer = self._new_gen_timer(job_id)
194 self._assert_model_init()
195 self._assert_args(height, width)
197 self.running = True # Mark running to avoid concurrent calls
199 try:
200 logging.info(
201 f"[{self.rank}] Generating image with {width}x{height}, "
202 f"{cfg:.1f} CFG, and "
203 f"{sampling_steps} steps, and "
204 f"'{prompt[:self.MAX_LOG_TEXT_LEN]}'.")
206 def callback_gen_timer(
207 pipeline: DiffusionPipeline,
208 step: int,
209 timestep: int,
210 callback_kwargs: Dict[str, Any],
211 ) -> Dict[str, Any]:
212 gen_timer.end(f"step_{step:03d}")
213 logging.info(f"[{self.rank}] Step {step + 1}/{sampling_steps}.")
215 if step < sampling_steps - 1:
216 gen_timer.start(f"step_{step + 1:03d}")
217 self.check_interrupted()
218 return callback_kwargs
220 image = await asyncio.to_thread(
221 self.model.generate_image, # type: ignore[union-attr]
222 prompt=prompt,
223 # image_size="auto",
224 # image_size=f"{height}x{width}",
225 image_size=(height, width),
226 diff_infer_steps=sampling_steps,
227 # diff_guidance_scale=cfg, # TODO figure if this breaks
228 seed=seed, # TODO figure if this breaks
229 stream=True,
230 )
231 """
232 # Pipeline mode would allow callback_on_step_end
233 cot_text = None # change for "think", "recaption"
234 system_prompt = None # get_system_prompt(use_system_prompt, bot_task, system_prompt)
235 model_inputs = self.model.prepare_model_inputs(
236 prompt=prompt,
237 cot_text=cot_text,
238 system_prompt=system_prompt,
239 mode="gen_image",
240 seed=seed,
241 image_size=image_size,
242 )
243 images = await asyncio.to_thread(
244 self.pipeline,
245 batch_size=1,
246 image_size=[width, height],
247 prompt=prompt,
248 num_inference_steps=sampling_steps,
249 guidance_scale=cfg,
250 output_type="pil", # TODO different output modes
251 callback_on_step_end=callback_gen_timer,
252 seed=seed,
253 **model_inputs
254 )
255 if not images or len(images) != 1:
256 raise RuntimeError(f"Wrong image generated: {images}")
257 image = images[0]
258 """
259 return image
260 finally:
261 self.running = False
262 gen_timer.end("total")
264 async def get_rest_args(self, data_json: Dict[str, str]) -> Dict[str, Any]:
265 if data_json is None:
266 raise ValueError("Missing JSON body")
267 prompt = data_json.get("prompt", None)
268 if prompt is None:
269 raise ValueError("Missing 'prompt' parameter")
270 height = int(data_json.get("height", 480))
271 width = int(data_json.get("width", 640))
272 steps = int(data_json.get("sampling_steps", 25))
273 seed = data_json.get("seed", None)
274 return {
275 "task": self.model_name,
276 "args": {
277 "prompt": prompt,
278 "width": width,
279 "height": height,
280 "sampling_steps": steps,
281 "seed": seed,
282 }
283 }