Added an example model with docs so people and agents can add models easier.
This commit is contained in:
Jaret Burkett
@@ -0,0 +1,158 @@
|
||||
"""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]
|
||||
Reference in New Issue
Block a user