Spaces:
Running on Zero
Running on Zero
File size: 6,527 Bytes
855e1f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """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]
|