Coverage for wrapper/qwenimage/wrapper_qwenimage.py: 89%
115 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
15from torch import inference_mode
17from wrapper_model import ModelGeneration
19from diffusers import QwenImagePipeline
21from xfuser.config import EngineConfig
24class QwenImageGeneration(ModelGeneration):
25 """Handle image generation using the QwenImage model."""
27 HF_MODEL_NAME = "Qwen/Qwen-Image"
29 def __init__(
30 self,
31 model_name: str = "qwenimage",
32 engine_config: EngineConfig = None,
33 param_dtype: torch.dtype = torch.bfloat16,
34 ) -> None:
35 super().__init__(model_name)
37 self.engine_config = engine_config
38 if self.engine_config is not None:
39 self.torch_compile = self.engine_config.runtime_config.use_torch_compile
40 self.param_dtype = param_dtype
42 # Model components
43 self.pipeline: Optional[QwenImagePipeline] = None
45 def __del__(self) -> None:
46 """Cleanup resources on deletion."""
47 if self.pipeline is not None:
48 self.pipeline = None
50 def init_parallelism(self) -> None:
51 """Initialize distributed settings for multi-GPU setups."""
52 self.load_timer.start("torch_dist")
54 self.rank = int(os.getenv("RANK", 0))
55 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
56 self.world_size = int(os.getenv("WORLD_SIZE", 1))
58 self.device_id = self.local_rank
59 self.device = torch.device(f"cuda:{self.device_id}")
61 torch.cuda.set_device(self.local_rank)
63 if self.world_size > 1:
64 logging.warning("Qwen Image is not optimized for multi-GPU setups (yet).")
65 self.world_size = 1
67 self.load_timer.end("torch_dist")
69 def load_model(self) -> None:
70 """Load the Qwen Image model from Hugging Face."""
71 assert torch.cuda.is_available()
73 self.load_timer.start("pipeline")
74 self.pipeline = QwenImagePipeline.from_pretrained(
75 pretrained_model_name_or_path=self.HF_MODEL_NAME,
76 torch_dtype=self.param_dtype,
77 )
78 self.pipeline = self.pipeline.to(self.device) # type: ignore[union-attr]
79 self.load_timer.end("pipeline")
81 def init_model_parallelism(self) -> None:
82 """Qwen Image does not support parallelism yet."""
83 if self.world_size > 1:
84 logging.warning("Parallelism not supported.")
86 def model_compile(self) -> None:
87 """Compile the model using torch.compile if enabled."""
88 if not self.torch_compile:
89 return
91 self.load_timer.start("dit_compile")
92 torch._inductor.config.reorder_for_compute_comm_overlap = True
93 assert self.pipeline is not None
94 self.pipeline.transformer = torch.compile( # type: ignore[attr-defined]
95 self.pipeline.transformer, # type: ignore[attr-defined]
96 mode="max-autotune-no-cudagraphs"
97 )
98 self.load_timer.end("dit_compile")
100 def _assert_model_init(self) -> None:
101 """Check if the model has been initialized."""
102 super()._assert_model_init()
103 assert self.pipeline is not None
105 def _get_vae_scale_factor(self) -> int:
106 """Return the VAE scale factor from the pipeline, raising if unavailable."""
107 if not self.pipeline:
108 raise ValueError("Pipeline not initialized.")
109 vae_scale_factor = getattr(self.pipeline, "vae_scale_factor", None)
110 if vae_scale_factor is None:
111 raise ValueError("Pipeline does not have vae_scale_factor.")
112 return vae_scale_factor
114 def _assert_args(
115 self,
116 height: int,
117 width: int,
118 ) -> None:
119 """Check if the image size is supported for the current parallelism setting."""
120 vae_scale_factor = self._get_vae_scale_factor()
121 height_latent = height // vae_scale_factor
122 width_latent = width // vae_scale_factor
123 img_latent_shape = (height_latent // 2) * (width_latent // 2)
124 if img_latent_shape % self.world_size != 0:
125 raise ValueError(f"{height}x{width} not supported for {self.world_size} GPUs.")
127 @inference_mode()
128 async def warmup(self) -> None:
129 """Warmup the Qwen Image model with a sample generation."""
130 logging.info(f"[{self.rank}] Warmup for Qwen Image generation.")
131 await self.generate(
132 # Ideally, we would use smaller sizes, but it has issues with 8 GPUs
133 width=1280,
134 height=800,
135 prompt="A warmup image to initialize the model.",
136 neg_prompt="",
137 sampling_steps=5) # It needs at least 5 steps to warm up properly
139 @override
140 @inference_mode()
141 async def generate(
142 self,
143 height: int,
144 width: int,
145 prompt: str,
146 neg_prompt: str = "",
147 sampling_steps: int = 25, # 10
148 seed: Optional[int] = None,
149 job_id: Optional[str] = None,
150 ) -> Image.Image:
151 """Generate an image using QwenImage."""
152 gen_timer = self._new_gen_timer(job_id)
154 self._assert_model_init()
155 self._assert_args(height, width)
156 assert self.pipeline is not None
158 self.running = True # Mark running to avoid concurrent calls
160 try:
161 if seed is None or seed < 0:
162 seed = random.randint(0, sys.maxsize)
163 seed_g = torch.Generator(device=self.device)
164 seed_g.manual_seed(seed)
166 def callback_gen_timer(
167 pipeline: QwenImagePipeline,
168 step: int,
169 timestep: int,
170 callback_kwargs: dict
171 ) -> dict:
172 gen_timer.end(f"step_{step:03d}")
173 if step < sampling_steps - 1:
174 gen_timer.start(f"step_{step + 1:03d}")
175 self.check_interrupted()
176 return callback_kwargs
178 gen_timer.start(f"step_{0:03d}")
179 output = self.pipeline( # type: ignore[operator]
180 height=height,
181 width=width,
182 prompt=prompt,
183 negative_prompt=neg_prompt,
184 num_inference_steps=sampling_steps,
185 output_type="pil",
186 generator=seed_g,
187 callback_on_step_end=callback_gen_timer,
188 )
190 if not output or len(output.images) != 1:
191 raise ValueError(f"Expected 1 image, but got {len(output.images)} images")
192 image = output.images[0]
193 return image
194 finally:
195 self.running = False
196 gen_timer.end("total")
198 async def get_rest_args(
199 self,
200 data_json: Dict[str, Union[str, int, float]]
201 ) -> Dict[str, Any]:
202 """Extract and validate REST API arguments for Qwen Image generation."""
203 if data_json is None:
204 raise ValueError("Missing JSON body")
205 prompt = data_json.get("prompt", None)
206 if prompt is None:
207 raise ValueError("Missing 'prompt' parameter")
208 neg_prompt = data_json.get("neg_prompt", "")
209 height = int(data_json.get("height", 480))
210 width = int(data_json.get("width", 640))
211 steps = int(data_json.get("sampling_steps", 20))
212 return {
213 "task": self.model_name,
214 "args": {
215 "prompt": prompt,
216 "neg_prompt": neg_prompt,
217 "width": width,
218 "height": height,
219 "sampling_steps": steps,
220 }
221 }