Coverage for wrapper/hidream/wrapper_hidream.py: 88%
141 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1import logging
2import os
3import sys
4import random
6from typing import override
7from typing import Optional
8from typing import Dict
9from typing import Any
10from typing import Union
12from PIL import Image
14import torch
15import torch.distributed as dist
16from torch import inference_mode
18from wrapper_model import ModelGeneration
20from diffusers import HiDreamImagePipeline
22from transformers import AutoTokenizer
23from transformers import LlamaForCausalLM
25from xfuser.config import EngineConfig
28class HiDreamGeneration(ModelGeneration):
29 """Handle image generation using the HiDream model."""
30 HF_MODEL_NAME = "HiDream-ai/HiDream-I1-Full"
32 def __init__(
33 self,
34 model_name: str = "hidream",
35 engine_config: Optional[EngineConfig] = None,
36 param_dtype: torch.dtype = torch.bfloat16,
37 ) -> None:
38 super().__init__(model_name)
40 self.engine_config = engine_config
41 if self.engine_config is not None:
42 self.torch_compile = self.engine_config.runtime_config.use_torch_compile
43 self.param_dtype = param_dtype
45 self.gpu = torch.cuda.get_device_name(0)
47 # Model components
48 self.pipeline: Optional[HiDreamImagePipeline] = None
49 self.text_encoder: Optional[LlamaForCausalLM] = None
50 self.tokenizer: Optional[AutoTokenizer] = None
52 def __del__(self) -> None:
53 # Clean models
54 if hasattr(self, "pipeline") and self.pipeline is not None:
55 self.pipeline = None
56 if hasattr(self, "text_encoder") and self.text_encoder is not None:
57 self.text_encoder = None
58 if hasattr(self, "tokenizer") and self.tokenizer is not None:
59 self.tokenizer = None
60 if dist.is_initialized():
61 dist.destroy_process_group()
63 def init_parallelism(self) -> None:
64 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))
70 self.device_id = self.local_rank
71 self.device = torch.device(f"cuda:{self.device_id}")
73 torch.cuda.set_device(self.local_rank)
75 if self.world_size > 1:
76 # https://github.com/dw763j/HiDream-I1-multigpu
77 # https://github.com/HiDream-ai/HiDream-I1/pull/30/files
78 logging.warning("HiDream is not optimized for multi-GPU setups (yet).")
79 self.world_size = 1
81 self.load_timer.end("torch_dist")
83 def load_model(self) -> None:
84 assert torch.cuda.is_available()
86 self.load_timer.start("pipeline")
88 self.tokenizer = AutoTokenizer.from_pretrained( # type: ignore[assignment]
89 "meta-llama/Meta-Llama-3.1-8B-Instruct") # nosec B615
90 assert self.tokenizer is not None
91 self.text_encoder = LlamaForCausalLM.from_pretrained(
92 "meta-llama/Meta-Llama-3.1-8B-Instruct",
93 output_hidden_states=True,
94 output_attentions=True,
95 torch_dtype=self.param_dtype,
96 ) # nosec B615
97 assert self.text_encoder is not None
99 self.pipeline = HiDreamImagePipeline.from_pretrained(
100 pretrained_model_name_or_path=self.HF_MODEL_NAME,
101 tokenizer_4=self.tokenizer,
102 text_encoder_4=self.text_encoder,
103 torch_dtype=self.param_dtype,
104 )
105 assert self.pipeline is not None
106 self.pipeline = self.pipeline.to(self.device) # type: ignore[union-attr]
107 assert self.pipeline is not None
108 self.load_timer.end("pipeline")
110 logging.info(
111 f"Loaded HiDreamImagePipeline: {self.HF_MODEL_NAME} device:{self.device} dtype:{self.param_dtype} "
112 f"device_map:{self.pipeline.hf_device_map}.") # type: ignore[union-attr]
114 def init_model_parallelism(self) -> None:
115 """HiDream does not support parallelism yet."""
116 if self.world_size > 1:
117 logging.warning("Parallelism not supported.")
119 def model_compile(self) -> None:
120 """Compile the model using torch.compile if enabled."""
121 if not self.torch_compile:
122 return
124 self.load_timer.start("dit_compile")
125 torch._inductor.config.reorder_for_compute_comm_overlap = True
126 assert self.pipeline is not None
127 assert self.pipeline.transformer is not None # type: ignore[attr-defined]
128 self.pipeline.transformer = torch.compile( # type: ignore[attr-defined]
129 self.pipeline.transformer, # type: ignore[attr-defined]
130 mode="max-autotune-no-cudagraphs"
131 )
132 assert self.pipeline.transformer is not None # type: ignore[attr-defined]
133 self.load_timer.end("dit_compile")
135 def _assert_model_init(self) -> None:
136 super()._assert_model_init()
137 assert self.pipeline is not None
139 def _get_vae_scale_factor(self) -> int:
140 """Return the VAE scale factor from the pipeline, raising if unavailable."""
141 if not self.pipeline:
142 raise ValueError("Pipeline not initialized.")
143 vae_scale_factor = getattr(self.pipeline, "vae_scale_factor", None)
144 if vae_scale_factor is None:
145 raise ValueError("Pipeline does not have vae_scale_factor.")
146 return vae_scale_factor
148 def _assert_args(
149 self,
150 height: int,
151 width: int,
152 ) -> None:
153 """Check if the image size is supported for the current parallelism setting."""
154 vae_scale_factor = self._get_vae_scale_factor()
155 height_latent = height // vae_scale_factor
156 width_latent = width // vae_scale_factor
157 img_latent_shape = (height_latent // 2) * (width_latent // 2)
158 if img_latent_shape % self.world_size != 0:
159 raise ValueError(f"{height}x{width} not supported for {self.world_size} GPUs.")
161 @inference_mode()
162 async def warmup(self) -> None:
163 logging.info(f"[{self.rank}] Warmup for HiDream generation.")
164 await self.generate(
165 height=720,
166 width=1280,
167 prompt="A warmup image to initialize the model.",
168 neg_prompt="",
169 sampling_steps=2)
171 @override
172 @inference_mode()
173 async def generate(
174 self,
175 height: int,
176 width: int,
177 prompt: str,
178 neg_prompt: str = "",
179 sampling_steps: int = 25, # 10\
180 seed: Optional[int] = None,
181 job_id: Optional[str] = None,
182 ) -> Image.Image:
183 """
184 Generate an image from a prompt using the HiDream model.
185 Args:
186 height (int): Height of the generated image.
187 width (int): Width of the generated image.
188 prompt (str): Text prompt to guide the image generation.
189 negative_prompt (str, optional): Negative prompt to avoid certain features in the image.
190 sampling_steps (int, optional): Number of inference steps for sampling. Default is 25.
191 seed (int, optional): Random seed for reproducibility. If None, a random seed will be generated.
192 job_id (str, optional): Job ID for tracking the generation process.
193 """
194 gen_timer = self._new_gen_timer(job_id)
196 self._assert_model_init()
197 self._assert_args(height, width)
198 assert self.pipeline is not None
200 self.running = True # Mark running to avoid concurrent calls
202 try:
203 if seed is None or seed < 0:
204 seed = random.randint(0, sys.maxsize)
205 seed_g = torch.Generator(device=self.device)
206 seed_g.manual_seed(seed)
208 def callback_gen_timer(
209 pipeline: HiDreamImagePipeline,
210 step: int,
211 timestep: int,
212 callback_kwargs: dict
213 ) -> dict:
214 gen_timer.end(f"step_{step:03d}")
215 if step < sampling_steps - 1:
216 gen_timer.start(f"step_{step + 1:03d}")
217 self.check_interrupted()
218 return callback_kwargs
220 gen_timer.start(f"step_{0:03d}")
221 output = self.pipeline( # type: ignore[operator]
222 height=height,
223 width=width,
224 prompt=prompt,
225 negative_prompt=neg_prompt,
226 num_inference_steps=sampling_steps,
227 output_type="pil",
228 generator=seed_g,
229 callback_on_step_end=callback_gen_timer,
230 )
232 if not output or len(output.images) != 1:
233 raise ValueError(f"Expected 1 image, but got {len(output.images)} images")
234 image = output.images[0]
235 return image
236 finally:
237 self.running = False
238 gen_timer.end("total")
240 def get_health(self) -> Dict[str, Any]:
241 ret = super().get_health()
242 ret.update({
243 "gpu": self.gpu,
244 "rank": self.rank,
245 "world_size": self.world_size,
246 "torch_compile": self.torch_compile,
247 "dtype": str(self.param_dtype),
248 "device_map": self.pipeline.hf_device_map if self.pipeline else None, # type: ignore[attr-defined]
249 })
250 return ret
252 async def get_rest_args(
253 self,
254 data_json: Dict[str, Union[str, int, float]]
255 ) -> Dict[str, Any]:
256 if data_json is None:
257 raise ValueError("Missing JSON body")
258 prompt = data_json.get("prompt", None)
259 if prompt is None:
260 raise ValueError("Missing 'prompt' parameter")
261 neg_prompt = data_json.get("neg_prompt", "")
262 height = int(data_json.get("height", 480))
263 width = int(data_json.get("width", 640))
264 steps = int(data_json.get("sampling_steps", 20))
265 seed = data_json.get("seed", None)
266 return {
267 "task": self.model_name,
268 "args": {
269 "prompt": prompt,
270 "neg_prompt": neg_prompt,
271 "height": height,
272 "width": width,
273 "sampling_steps": steps,
274 "seed": seed,
275 }
276 }