Coverage for wrapper/llamagen/wrapper_llamagen.py: 65%
211 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Based on:
3https://github.com/FoundationVision/LlamaGen/blob/main/autoregressive/sample/sample_t2i.py
4"""
6import logging
7import os
8import sys
9import random
10import numpy as np
12from typing import override
13from typing import Optional
14from typing import Dict
15from typing import List
16from typing import Tuple
17from typing import Any
18from typing import Union
20from PIL import Image
22import torch
23import torch.distributed as dist
24from torch import inference_mode
26from wrapper_model import ModelGeneration
28from tokenizer.tokenizer_image.vq_model import VQ_models
29from language.t5 import T5Embedder
30from autoregressive.models.gpt import GPT_models
31from autoregressive.models.generate import generate
33from xfuser.config import EngineConfig
36class LlamaGenGeneration(ModelGeneration):
37 """LlamaGen model for text-to-image generation."""
39 def __init__(
40 self,
41 model_name: str = "llamagen",
42 engine_config: EngineConfig = None,
43 param_dtype: torch.dtype = torch.bfloat16,
44 ) -> None:
45 super().__init__(model_name)
47 self.engine_config = engine_config
48 self.torch_compile = True
49 if self.engine_config is not None:
50 self.torch_compile = self.engine_config.runtime_config.use_torch_compile
51 self.param_dtype = param_dtype
53 # LlamaGen specific configuration
54 self.REPO_ID = "peizesun/llamagen_t2i"
55 self.GPT_TYPE = "t2i"
56 self.VQ_MODEL_NAME = "VQ-16"
57 # GPT-B 111M
58 # GPT-L 343M
59 # GPT-XL 775M
60 # GPT-3B 3.1B
61 self.GPT_MODEL_NAME = "GPT-XL"
63 # LlamaGen specific parameters
64 self.downsample_size = 16
65 self.codebook_size = 16384
66 self.codebook_embed_dim = 8
67 self.cls_token_num = 120
69 # T5 model configuration - match this with your trained GPT model
70 # Common T5 models and their dimensions:
71 # - flan-t5-xl: 2048 dimensions
72 # - t5-v1_1-xl: 2048 dimensions
73 # - t5-v1_1-xxl: 4096 dimensions
74 # - flan-t5-xxl: 4096 dimensions
75 self.T5_MODEL_TYPE = "flan-t5-xl"
76 self.t5_feature_max_len = 120
77 self.t5_feature_dim = 2048 # This should match the T5 model's hidden size
78 # It supports 256 and 512
79 self.image_size = 512 # pixels
80 self.latent_size = self.image_size // self.downsample_size
82 # Parallelism
83 self.gpu: Optional[str] = None
84 if torch.cuda.is_available():
85 self.gpu = torch.cuda.get_device_name(0)
87 # Model components
88 self.vq_model: Optional[torch.nn.Module] = None
89 self.gpt_model: Optional[torch.nn.Module] = None
90 self.t5_model: Optional[T5Embedder] = None
92 def __del__(self) -> None:
93 if self.vq_model is not None:
94 self.vq_model = None
95 if self.gpt_model is not None:
96 self.gpt_model = None
97 if self.t5_model is not None:
98 self.t5_model = None
99 if dist.is_initialized():
100 dist.destroy_process_group()
102 def init_parallelism(self) -> None:
103 self.load_timer.start("torch_dist")
105 self.rank = int(os.getenv("RANK", 0))
106 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
107 self.world_size = int(os.getenv("WORLD_SIZE", 1))
109 self.device_id = self.local_rank
110 self.device = torch.device(f"cuda:{self.device_id}")
112 torch.cuda.set_device(self.local_rank)
114 if self.world_size > 1:
115 logging.warning("LlamaGen is not optimized for multi-GPU setups (yet).")
116 self.world_size = 1
118 self.load_timer.end("torch_dist")
120 def load_model(self) -> None:
121 assert torch.cuda.is_available()
123 # Setup PyTorch optimizations like in the original code
124 torch.backends.cuda.matmul.allow_tf32 = True
125 torch.backends.cudnn.allow_tf32 = True
126 torch.set_float32_matmul_precision('high')
127 setattr(torch.nn.Linear, 'reset_parameters', lambda self: None)
128 setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None)
129 os.environ["TOKENIZERS_PARALLELISM"] = "false"
131 self.load_timer.start("vq_model")
132 # Load VQ model
133 self.vq_model = VQ_models[self.VQ_MODEL_NAME](
134 codebook_size=self.codebook_size,
135 codebook_embed_dim=self.codebook_embed_dim
136 )
137 self.vq_model.to(self.device)
138 self.vq_model.eval()
139 vq_file = f"{self.REPO_ID}/vq_ds16_t2i.pt"
140 checkpoint = torch.load(
141 vq_file,
142 map_location="cpu",
143 weights_only=False) # nosec B614 - trusted HuggingFace model checkpoint
144 self.vq_model.load_state_dict(checkpoint["model"])
145 self.load_timer.end("vq_model")
147 self.load_timer.start("gpt_model")
148 self.gpt_model = GPT_models[self.GPT_MODEL_NAME](
149 block_size=self.latent_size ** 2,
150 cls_token_num=self.cls_token_num,
151 model_type=self.gpt_type,
152 )
153 assert self.gpt_model is not None
154 self.gpt_model = self.gpt_model.to(dtype=self.param_dtype)
155 self.gpt_model = self.gpt_model.to(device=self.device)
156 if self.image_size not in [256, 512]:
157 raise ValueError(f"Image size {self.image_size} not supported. Must be 256 or 512.")
158 t2i_file = f"{self.REPO_ID}/t2i_XL_stage2_{self.image_size}.pt"
159 checkpoint = torch.load(
160 t2i_file,
161 map_location="cpu",
162 weights_only=False) # nosec B614 - trusted HuggingFace model checkpoint
163 model_weight = checkpoint.get("model", checkpoint.get("module", checkpoint.get("state_dict", checkpoint)))
164 self.gpt_model.load_state_dict(model_weight, strict=False)
165 self.gpt_model.eval()
166 self.load_timer.end("gpt_model")
168 self.load_timer.start("t5_model")
169 # Load T5 model for text embeddings which should be already downloaded:
170 # For flan-t5-xl: huggingface-cli download google/flan-t5-xl --local-dir google/flan-t5-xl
171 # https://github.com/FoundationVision/LlamaGen/blob/main/language/t5.py
172 # self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
173 # self.model = T5EncoderModel.from_pretrained(path, **t5_model_kwargs).eval()
174 # t5_path = "DeepFloyd" # "t5-v1_1-xxl"
175 t5_path = "google" # "flan-t5-xl"
176 logging.info(f"Loading T5 model from '{t5_path}'.")
177 if not os.path.exists(t5_path):
178 raise FileNotFoundError(
179 f"T5 model directory '{t5_path}' does not exist. Please download the model using huggingface-cli.")
180 if self.T5_MODEL_TYPE != ["flan-t5-xl"]:
181 raise ValueError("Only 'flan-t5-xl' is available.")
182 self.t5_model = T5Embedder(
183 device=self.device,
184 local_cache=True,
185 cache_dir=t5_path if os.path.exists(t5_path) else None,
186 dir_or_name=self.T5_MODEL_TYPE,
187 torch_dtype=self.param_dtype,
188 model_max_length=self.t5_feature_max_len,
189 )
190 self.load_timer.end("t5_model")
192 logging.info(f"Loaded LlamaGen. VQ:{self.vq_model_name} GPT:{self.gpt_model_name} T5:{self.t5_model_type}.")
194 def init_model_parallelism(self) -> None:
195 if self.world_size > 1:
196 logging.warning("LlamaGen does not support model parallelism yet.")
198 def model_compile(self) -> None:
199 if not self.torch_compile:
200 return
202 self.load_timer.start("model_compile")
203 assert self.gpt_model is not None
204 self.gpt_model = torch.compile( # type: ignore[assignment]
205 self.gpt_model,
206 mode="reduce-overhead",
207 fullgraph=True
208 )
209 self.load_timer.end("model_compile")
211 def _assert_model_init(self) -> None:
212 assert self.vq_model is not None
213 assert self.gpt_model is not None
214 assert self.t5_model is not None
216 def _assert_args(self, image_size: int) -> None:
217 if image_size not in [256, 384, 512]:
218 raise ValueError(f"Image size {image_size} not supported. Must be 256, 384, or 512.")
219 if image_size != self.image_size:
220 raise ValueError(f"Image size {image_size} does not match model's size {self.image_size}.")
222 def _prepare_text_embeddings(
223 self,
224 prompts: List[str],
225 no_left_padding: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
226 """
227 Prepare text embeddings using T5 model.
228 """
229 assert self.t5_model is not None
230 caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)
232 if not no_left_padding:
233 # Implement left-padding as in the original code
234 new_emb_masks = torch.flip(emb_masks, dims=[-1])
235 new_caption_embs_list = []
236 for idx, (caption_emb, emb_mask) in enumerate(zip(caption_embs, emb_masks)):
237 valid_num = int(emb_mask.sum().item())
238 logging.debug(f'Prompt {idx} token len: {valid_num}')
239 new_caption_emb = torch.cat([caption_emb[valid_num:], caption_emb[:valid_num]])
240 new_caption_embs_list.append(new_caption_emb)
241 new_caption_embs = torch.stack(new_caption_embs_list)
242 else:
243 new_caption_embs, new_emb_masks = caption_embs, emb_masks
245 c_indices = new_caption_embs * new_emb_masks[:, :, None]
246 c_emb_masks = new_emb_masks
248 return c_indices, c_emb_masks
250 def _sample_to_image(self, sample: torch.Tensor) -> Image.Image:
251 # "sample" is a tensor with shape [C, H, W] [3, H, W] and values in [-1, 1]
252 sample = (sample + 1) / 2 # Convert from [-1, 1] to [0, 1]
253 sample = torch.clamp(sample, 0, 1)
254 # Convert to numpy and transpose to HWC format
255 sample_np = sample.cpu().numpy().transpose(1, 2, 0)
256 sample_np = (sample_np * 255).astype(np.uint8) # Convert to [0, 255] uint8
257 pil_image = Image.fromarray(sample_np)
258 return pil_image
260 @inference_mode()
261 async def warmup(self) -> None:
262 logging.info(f"[{self.rank}] Warmup for LlamaGen generation.")
263 await self.generate(
264 prompt="A warmup image to initialize the model.",
265 image_size=512,
266 cfg_scale=7.5,
267 temperature=1.0,
268 top_k=1000,
269 )
271 @override
272 @inference_mode()
273 async def generate(
274 self,
275 prompt: str,
276 image_size: int = 512, # pixels (512 x 512)
277 cfg_scale: float = 7.5,
278 temperature: float = 1.0,
279 top_k: int = 1000,
280 top_p: float = 1.0,
281 seed: Optional[int] = None,
282 no_left_padding: bool = False,
283 job_id: Optional[str] = None,
284 ) -> Image.Image:
285 """
286 Generate an image from a prompt using the LlamaGen model.
287 Args:
288 prompt (str): Text prompt to guide the image generation.
289 image_size (int): Size of the generated image (256, 384, or 512).
290 cfg_scale (float): Classifier-free guidance scale.
291 temperature (float): Sampling temperature.
292 top_k (int): Top-k sampling parameter.
293 top_p (float): Top-p (nucleus) sampling parameter.
294 seed (int): Random seed for reproducibility.
295 no_left_padding (bool): Whether to skip left padding for text embeddings.
296 Returns:
297 Image.Image: Generated PIL Image.
298 """
299 gen_timer = self._new_gen_timer(job_id)
301 self._assert_model_init()
302 self._assert_args(image_size)
304 self.running = True # Mark running to avoid concurrent calls
306 try:
307 # Set seed for reproducibility
308 if seed is not None:
309 torch.manual_seed(seed)
310 torch.cuda.manual_seed(seed)
311 else:
312 seed = random.randint(0, sys.maxsize)
313 torch.manual_seed(seed)
315 torch.set_grad_enabled(False)
316 latent_size = image_size // self.downsample_size
318 logging.info(f"Generating image with prompt: '{prompt}'.")
319 logging.info(f"Image size: {image_size}, Latent size: {latent_size}.")
321 # Prepare text embeddings
322 gen_timer.start("text_embeddings")
323 prompts = [prompt] # LlamaGen expects a list
324 c_indices, c_emb_masks = self._prepare_text_embeddings(prompts, no_left_padding)
325 gen_timer.end("text_embeddings")
327 # Generate token indices
328 gen_timer.start("sampling")
329 qzshape = [len(c_indices), self.codebook_embed_dim, latent_size, latent_size]
331 index_sample = generate(
332 self.gpt_model,
333 c_indices,
334 latent_size ** 2,
335 c_emb_masks,
336 cfg_scale=cfg_scale,
337 temperature=temperature,
338 top_k=top_k,
339 top_p=top_p,
340 sample_logits=True,
341 )
342 gen_timer.end("sampling")
344 # Decode to image
345 gen_timer.start("decoding")
346 assert self.vq_model is not None
347 # output in [-1, 1]
348 samples = self.vq_model.decode_code(index_sample, qzshape) # type: ignore [operator]
349 sample = samples[0]
350 gen_timer.end("decoding")
352 # Convert to PIL Image
353 gen_timer.start("convert_pil")
354 pil_image = self._sample_to_image(sample)
355 gen_timer.end("convert_pil")
357 return pil_image
358 finally:
359 self.running = False
360 gen_timer.end("total")
362 def get_health(self) -> Dict[str, Any]:
363 ret = super().get_health()
364 ret.update({
365 "gpu": self.gpu,
366 "rank": self.rank,
367 "world_size": self.world_size,
368 "torch_compile": self.torch_compile,
369 "dtype": str(self.param_dtype),
370 "vq_model": self.vq_model_name,
371 "gpt_model": self.gpt_model_name,
372 "gpt_type": self.gpt_type,
373 "t5_model": self.t5_model_type,
374 })
375 return ret
377 async def get_rest_args(
378 self,
379 data_json: Dict[str, Union[str, int, float]]
380 ) -> Dict[str, Any]:
381 if data_json is None or not isinstance(data_json, dict):
382 raise ValueError("Missing JSON body")
384 prompt = data_json.get("prompt", None)
385 if prompt is None:
386 raise ValueError("Missing 'prompt' parameter")
388 image_size = int(data_json.get("image_size", 512))
389 cfg_scale = float(data_json.get("cfg_scale", 7.5))
390 temperature = float(data_json.get("temperature", 1.0))
391 top_k = int(data_json.get("top_k", 1000))
392 top_p = float(data_json.get("top_p", 1.0))
393 seed = data_json.get("seed", None)
394 if seed is not None:
395 seed = int(seed)
396 no_left_padding = bool(data_json.get("no_left_padding", False))
398 return {
399 "task": self.model_name,
400 "args": {
401 "prompt": prompt,
402 "image_size": image_size,
403 "cfg_scale": cfg_scale,
404 "temperature": temperature,
405 "top_k": top_k,
406 "top_p": top_p,
407 "seed": seed,
408 "no_left_padding": no_left_padding,
409 }
410 }