Coverage for wrapper/cogview/wrapper_cogview.py: 45%
110 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2from diffusers import CogView4Pipeline
3import torch
5pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16).to("cuda")
7# Open it for reduce GPU memory usage
8pipe.enable_model_cpu_offload()
9pipe.vae.enable_slicing()
10pipe.vae.enable_tiling()
12prompt = "A vibrant cherry red sports car sits proudly under the gleaming sun, its polished exterior smooth "
13prompt += "and flawless, casting a mirror-like reflection. The car features a low, aerodynamic body, angular "
14prompt += "headlights that gaze forward like predatory eyes, and a set of black, high-gloss racing rims that "
15prompt += "contrast starkly with the red. "
16prompt += "A subtle hint of chrome embellishes the grille and exhaust, while the tinted windows suggest a "
17prompt += "luxurious and private interior. The scene conveys a sense of speed and elegance, the car appearing as "
18prompt += "if it's about to burst into a sprint along a coastal road, with the ocean's azure waves crashing in "
19prompt += "the background."
20image = pipe(
21 prompt=prompt,
22 guidance_scale=3.5,
23 num_images_per_prompt=1,
24 num_inference_steps=50,
25 width=1024,
26 height=1024,
27).images[0]
29image.save("cogview4.png")
30"""
32import logging
33import os
34import sys
35import random
37from typing import Optional
38from typing import Dict
39from typing import Union
40from typing import Any
42from PIL import Image
44import torch
45from torch import inference_mode
47from wrapper_model import ModelGeneration
49from diffusers import CogView4Pipeline
51from xfuser.config import EngineConfig
54class CogViewGeneration(ModelGeneration):
55 """Handle image generation using the CogView model."""
56 MODEL_NAME = "THUDM/CogView4-6B"
58 def __init__(
59 self,
60 model_name: str = "cogview",
61 engine_config: Optional[EngineConfig] = None,
62 param_dtype: torch.dtype = torch.bfloat16,
63 ) -> None:
64 super().__init__(model_name)
66 self.engine_config = engine_config
67 if self.engine_config is not None:
68 self.torch_compile = self.engine_config.runtime_config.use_torch_compile
69 self.param_dtype = param_dtype
71 # Parallelism
72 self.GPU = None
73 if torch.cuda.is_available():
74 self.GPU = torch.cuda.get_device_name(0)
76 self.base_seed = random.randint(0, sys.maxsize)
78 # Model components
79 self.pipeline: Optional[CogView4Pipeline] = None
81 def __del__(self) -> None:
82 # Clean models
83 if self.pipeline is not None:
84 self.pipeline = None
86 def init_parallelism(self) -> None:
87 self.load_timer.start("torch_dist")
89 self.rank = int(os.getenv("RANK", 0))
90 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
91 self.world_size = int(os.getenv("WORLD_SIZE", 1))
93 self.device_id = self.local_rank
94 self.device = torch.device(f"cuda:{self.device_id}")
96 torch.cuda.set_device(self.local_rank)
98 if self.world_size > 1:
99 logging.warning("CogView is not optimized for multi-GPU setups (yet).")
100 self.world_size = 1
102 self.load_timer.end("torch_dist")
104 def load_model(self) -> None:
105 assert torch.cuda.is_available()
107 self.load_timer.start("pipeline")
108 self.pipeline = CogView4Pipeline.from_pretrained(
109 pretrained_model_name_or_path=self.MODEL_NAME,
110 torch_dtype=self.param_dtype,
111 )
112 assert self.pipeline is not None
113 self.pipeline = self.pipeline.to(self.device) # type: ignore[union-attr]
115 # Enable memory optimizations
116 """
117 self.pipeline.enable_model_cpu_offload()
118 self.pipeline.vae.enable_slicing()
119 self.pipeline.vae.enable_tiling()
120 """
122 self.load_timer.end("pipeline")
124 logging.info(f"Loaded CogView4Pipeline: {self.MODEL_NAME} device:{self.device} dtype:{self.param_dtype}.")
126 def init_model_parallelism(self) -> None:
127 # CogView4 doesn't support model parallelism yet
128 pass
130 def model_compile(self) -> None:
131 if not self.torch_compile:
132 return
134 # TODO this is not likely supported
135 self.load_timer.start("compile")
136 torch._inductor.config.reorder_for_compute_comm_overlap = True
137 assert self.pipeline is not None
138 if hasattr(self.pipeline, 'transformer'):
139 self.pipeline.transformer = torch.compile(
140 self.pipeline.transformer,
141 mode="max-autotune-no-cudagraphs"
142 )
143 self.load_timer.end("compile")
145 def _assert_model_init(self) -> None:
146 super()._assert_model_init()
147 assert self.pipeline is not None
149 def _assert_args(
150 self,
151 height: int,
152 width: int,
153 ) -> None:
154 """
155 Check if the image size is supported for the current parallelism setting.
156 """
157 # CogView4 has specific size requirements
158 if height % 64 != 0 or width % 64 != 0:
159 raise ValueError(f"Height and width must be divisible by 64, got {height}x{width}")
161 @inference_mode()
162 async def warmup(self) -> None:
163 logging.info(f"[{self.rank}] Warmup for CogView4 generation.")
164 await self.generate(
165 height=512,
166 width=512,
167 prompt="A warmup image to initialize the model.",
168 sampling_steps=2
169 )
171 @inference_mode()
172 async def generate(
173 self,
174 height: int,
175 width: int,
176 prompt: str,
177 guidance_scale: float = 3.5,
178 sampling_steps: int = 50,
179 job_id: Optional[str] = None,
180 ) -> Image.Image:
181 """
182 Generate an image from a prompt using the CogView4 model.
183 Args:
184 height (int): Height of the generated image.
185 width (int): Width of the generated image.
186 prompt (str): Text prompt to guide the image generation.
187 guidance_scale (float, optional): Guidance scale for generation. Default is 3.5.
188 sampling_steps (int, optional): Number of inference steps for sampling. Default is 50.
189 """
190 gen_timer = self._new_gen_timer(job_id)
192 self._assert_model_init()
193 self._assert_args(height, width)
194 assert self.pipeline is not None
196 self.running = True # Mark running to avoid concurrent calls
198 try:
199 seed = self.base_seed if self.base_seed >= 0 else random.randint(0, sys.maxsize)
200 seed_g = torch.Generator(device=self.device)
201 seed_g.manual_seed(seed)
203 def callback_gen_timer(
204 pipeline: CogView4Pipeline,
205 step: int,
206 timestep: int,
207 callback_kwargs: dict
208 ) -> dict:
209 gen_timer.end(f"step_{step:03d}")
210 if step < sampling_steps - 1:
211 gen_timer.start(f"step_{step + 1:03d}")
212 return callback_kwargs
214 gen_timer.start(f"step_{0:03d}")
215 output = self.pipeline( # type: ignore[operator]
216 prompt=prompt,
217 height=height,
218 width=width,
219 guidance_scale=guidance_scale,
220 num_images_per_prompt=1,
221 num_inference_steps=sampling_steps,
222 output_type="pil",
223 generator=seed_g,
224 callback_on_step_end=callback_gen_timer,
225 )
226 images = output.images
228 assert len(images) == 1, f"Expected 1 image, but got {len(images)} images."
230 return images[0]
231 finally:
232 self.running = False
233 gen_timer.end("total")
235 def get_health(self) -> Dict[str, Any]:
236 ret = super().get_health()
237 ret.update({
238 "gpu": self.GPU,
239 "rank": self.rank,
240 "world_size": self.world_size,
241 "torch_compile": self.torch_compile,
242 "dtype": str(self.param_dtype),
243 "device_map": getattr(self.pipeline, 'hf_device_map', None) if self.pipeline else None,
244 })
245 return ret
247 async def get_rest_args(
248 self,
249 data_json: Dict[str, Union[str, int, float]]
250 ) -> Dict[str, Any]:
251 if data_json is None or not isinstance(data_json, dict):
252 raise ValueError("Missing JSON body")
254 prompt = data_json.get("prompt", None)
255 if prompt is None:
256 raise ValueError("Missing 'prompt' parameter")
258 height = int(data_json.get("height", 1024))
259 width = int(data_json.get("width", 1024))
260 guidance_scale = float(data_json.get("guidance_scale", 3.5))
261 steps = int(data_json.get("sampling_steps", 50))
262 return {
263 "task": self.model_name,
264 "args": {
265 "prompt": prompt,
266 "height": height,
267 "width": width,
268 "guidance_scale": guidance_scale,
269 "sampling_steps": steps,
270 }
271 }