Generate predictions during validation with Advanced Classifier-Free Guidance.
(
model,
prompt_emb,
prompt_emb_null,
latents,
latents_mask,
ref_latents,
ref_latents_mask,
num_inference_steps: int = 50,
cfg_scale: float = 5.0,
use_ema: bool = False,
device: torch.device = torch.device('cuda'),
dtype: torch.dtype = torch.bfloat16,
scheduler: FlowMatchScheduler = None,
seed: int = None,
logger=None,
condition_on_text: bool = False,
attend_to_text_mask: torch.Tensor | None = None,
)
| 277 | |
| 278 | @torch.no_grad() |
| 279 | def generate_pipe( |
| 280 | model, |
| 281 | prompt_emb, |
| 282 | prompt_emb_null, |
| 283 | latents, |
| 284 | latents_mask, |
| 285 | ref_latents, |
| 286 | ref_latents_mask, |
| 287 | num_inference_steps: int = 50, |
| 288 | cfg_scale: float = 5.0, |
| 289 | use_ema: bool = False, |
| 290 | device: torch.device = torch.device('cuda'), |
| 291 | dtype: torch.dtype = torch.bfloat16, |
| 292 | scheduler: FlowMatchScheduler = None, |
| 293 | seed: int = None, |
| 294 | logger=None, |
| 295 | condition_on_text: bool = False, |
| 296 | attend_to_text_mask: torch.Tensor | None = None, |
| 297 | ): |
| 298 | """Generate predictions during validation with Advanced Classifier-Free Guidance.""" |
| 299 | to_eval_mode() |
| 300 | generator = torch.Generator(device).manual_seed(seed if seed is not None else torch.randint(0, 1000000, (1,)).item()) |
| 301 | |
| 302 | # Use EMA model if specified |
| 303 | inf_model = ema if use_ema else model |
| 304 | |
| 305 | # Prepare noise and initial latents |
| 306 | noise = randn_tensor( |
| 307 | logger, |
| 308 | latents.shape, |
| 309 | generator=generator, |
| 310 | device=device, |
| 311 | dtype=dtype |
| 312 | ) |
| 313 | |
| 314 | # Set up scheduler for inference |
| 315 | scheduler.set_timesteps(num_inference_steps, training=False, denoising_strength=0.7) |
| 316 | timesteps = scheduler.timesteps.to(device) |
| 317 | xt = noise # Start with pure noise for generation |
| 318 | |
| 319 | # Pad prompt_emb_null to the same length as prompt_emb # [B, L, C] |
| 320 | # prompt_emb_null: [B, L1, C], prompt_emb: [B, L2, C] |
| 321 | if prompt_emb_null.size(1) < prompt_emb.size(1): |
| 322 | prompt_emb_zeros = torch.zeros(prompt_emb.size(0), prompt_emb.size(1) - prompt_emb_null.size(1), prompt_emb.size(2), device=prompt_emb.device, dtype=prompt_emb.dtype) |
| 323 | prompt_emb_null = torch.cat([prompt_emb_null, prompt_emb_zeros], dim=1) |
| 324 | |
| 325 | # Denoising loop with Advanced CFG |
| 326 | latents_mask_input = torch.cat([latents_mask] * 2, dim=0) |
| 327 | ref_latents_null = torch.zeros_like(ref_latents) |
| 328 | ref_latents_input = torch.cat([ref_latents, ref_latents_null], dim=0) |
| 329 | ref_latents_mask_input = torch.cat([ref_latents_mask] * 2, dim=0) |
| 330 | attend_to_text_mask_input = None |
| 331 | if attend_to_text_mask is not None: |
| 332 | attend_to_text_mask_input = torch.cat([attend_to_text_mask] * 2, dim=0) |
| 333 | |
| 334 | # Contexts |
| 335 | context_input = torch.cat([ |
| 336 | prompt_emb, # Conditional |
no test coverage detected