159 lines
6.4 KiB
Python
159 lines
6.4 KiB
Python
"""A minimal sampling pipeline for the example model.
|
|
|
|
ai-toolkit only uses your pipeline to render preview/sample images during
|
|
training (see BaseModel.generate_images -> ExampleModel.generate_single_image).
|
|
It does NOT need to be a diffusers DiffusionPipeline, and because ai-toolkit
|
|
always encodes the prompts itself (so it can cache embeds, apply trigger words,
|
|
run adapters, etc.) the pipeline never sees raw prompt strings -- only
|
|
already-encoded ``AdvancedPromptEmbeds``.
|
|
|
|
So all a pipeline has to do is:
|
|
|
|
1. make starting noise
|
|
2. loop the scheduler over timesteps, calling the transformer
|
|
3. apply classifier-free guidance (cond vs uncond prediction)
|
|
4. decode the final latents with the VAE and return PIL images
|
|
|
|
The pattern of passing the whole BaseModel instance into the pipeline (rather
|
|
than individual components) is borrowed from ../../ideogram4/src/pipeline.py.
|
|
It keeps the pipeline tiny because it can reuse the model's scheduler factory,
|
|
``decode_latents`` and device/dtype bookkeeping.
|
|
"""
|
|
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
|
|
|
|
def pad_prompt_embeds(
|
|
embeds_list: List[torch.Tensor],
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
):
|
|
"""Right-pad a list of per-sample text features into one batch tensor.
|
|
|
|
in: embeds_list list (len B) of (L_i, D) tensors -- this is exactly what
|
|
``AdvancedPromptEmbeds.text_embeds`` holds: one tensor per
|
|
batch item, each at its own natural length.
|
|
out: features (B, L_max, D) zero-padded on the right
|
|
mask (B, L_max) long, 1 = real token, 0 = padding
|
|
|
|
Storing embeds unpadded per item and only padding at the model call is the
|
|
preferred pattern: cached embeds stay small, and items of very different
|
|
prompt lengths can share a batch.
|
|
"""
|
|
lengths = [e.shape[0] for e in embeds_list]
|
|
max_len = max(lengths)
|
|
dim = embeds_list[0].shape[-1]
|
|
batch_size = len(embeds_list)
|
|
|
|
features = torch.zeros(batch_size, max_len, dim, device=device, dtype=dtype)
|
|
mask = torch.zeros(batch_size, max_len, dtype=torch.long, device=device)
|
|
for i, e in enumerate(embeds_list):
|
|
n = e.shape[0]
|
|
features[i, :n] = e.to(device, dtype)
|
|
mask[i, :n] = 1
|
|
return features, mask
|
|
|
|
|
|
class ExamplePipeline:
|
|
"""Lightweight flow-matching sampler used for training previews."""
|
|
|
|
def __init__(self, model):
|
|
# ``model`` is the ExampleModel (a BaseModel subclass), giving us
|
|
# access to model.transformer, model.vae, model.decode_latents, etc.
|
|
self.model = model
|
|
|
|
@property
|
|
def device(self):
|
|
return self.model.device_torch
|
|
|
|
def to(self, *args, **kwargs):
|
|
# BaseModel.generate_images may call pipeline.to(device); we manage
|
|
# devices through the model itself, so this is a no-op.
|
|
return self
|
|
|
|
def set_progress_bar_config(self, **kwargs):
|
|
# called by the sampler harness (inside a try/except, so optional);
|
|
# diffusers pipelines use it to silence tqdm. Nothing to do here.
|
|
pass
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
# AdvancedPromptEmbeds with key ``text_embeds`` (list of (L, D) tensors)
|
|
conditional_embeds,
|
|
unconditional_embeds,
|
|
height: int = 1024,
|
|
width: int = 1024,
|
|
num_inference_steps: int = 25,
|
|
guidance_scale: float = 4.0,
|
|
latents: Optional[torch.Tensor] = None, # pre-made noise, usually None
|
|
generator: Optional[torch.Generator] = None, # seeded RNG for reproducible samples
|
|
**kwargs,
|
|
) -> List[Image.Image]:
|
|
model = self.model
|
|
device = model.device_torch
|
|
dtype = model.torch_dtype
|
|
transformer = model.transformer
|
|
|
|
# Always sample with a FRESH scheduler. The training scheduler is
|
|
# stateful; mutating it mid-training would corrupt the train step.
|
|
scheduler = model.get_train_scheduler()
|
|
scheduler.set_timesteps(num_inference_steps, device=device)
|
|
timesteps = scheduler.timesteps # 1000 -> 0 scale
|
|
|
|
# pixel size -> latent size (VAE downsample only; the transformer
|
|
# patchifies internally so latents stay unpacked here)
|
|
gh = height // model.vae_scale_factor
|
|
gw = width // model.vae_scale_factor
|
|
|
|
do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
|
|
|
|
# 1. starting noise (keep it float32; cast per model call)
|
|
if latents is None:
|
|
shape = (1, transformer.in_channels, gh, gw)
|
|
latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
|
|
latents = latents.to(device, dtype=torch.float32)
|
|
|
|
# 2. pad the per-item embed lists into batch tensors once, up front
|
|
cond_feats, cond_mask = pad_prompt_embeds(conditional_embeds.text_embeds, device, dtype)
|
|
if do_cfg:
|
|
uncond_feats, uncond_mask = pad_prompt_embeds(unconditional_embeds.text_embeds, device, dtype)
|
|
|
|
# 3. denoising loop
|
|
for t in timesteps:
|
|
# scheduler timesteps are on a 0-1000 scale; the transformer wants
|
|
# flow time in [0, 1] with 1 = pure noise
|
|
t01 = (t / 1000.0).to(device).expand(latents.shape[0])
|
|
|
|
v_cond = transformer(
|
|
hidden_states=latents.to(dtype),
|
|
timestep=t01,
|
|
encoder_hidden_states=cond_feats,
|
|
attention_mask=cond_mask,
|
|
)
|
|
if do_cfg:
|
|
v_uncond = transformer(
|
|
hidden_states=latents.to(dtype),
|
|
timestep=t01,
|
|
encoder_hidden_states=uncond_feats,
|
|
attention_mask=uncond_mask,
|
|
)
|
|
# classifier-free guidance: push the prediction away from the
|
|
# unconditional (negative prompt) direction
|
|
v = v_uncond + guidance_scale * (v_cond - v_uncond)
|
|
else:
|
|
v = v_cond
|
|
|
|
latents = scheduler.step(v.to(torch.float32), t, latents, return_dict=False)[0]
|
|
|
|
# 4. decode latents -> images in [-1, 1] -> uint8 PIL
|
|
images = model.decode_latents(latents, device=device, dtype=dtype)
|
|
images = images.float().clamp(-1.0, 1.0)
|
|
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
|
|
images = images.permute(0, 2, 3, 1).cpu().numpy()
|
|
return [Image.fromarray(arr) for arr in images]
|