Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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 | |
| 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] | |