| """Inference: multi-condition classifier-free guidance + ODE (docs/04 "Çok koşullu CFG"). |
| |
| Three passes per step, batched together: |
| |
| * ``v∅`` — unconditional: text dropped, emotion/intensity null (the prompt policy is a |
| parameter: ``drop_prompt_uncond``); |
| * ``vT`` — text + tags, emotion null; |
| * ``vTE`` — text + tags + emotion. |
| |
| ``v = v∅ + w_text · (vT − v∅) + w_emo · (vTE − vT)``. With ``w_text = w_emo = 1`` this is |
| exactly ``vTE``. ``w_emo`` is the user-facing "emotion intensity" knob. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import replace |
| from typing import Literal |
|
|
| import torch |
| from torch import Tensor |
|
|
| from nida.model.dit import Conditions, DiT |
| from nida.model.flow import sample_ode |
|
|
| __all__ = ["guided_velocity", "synthesize_latents"] |
|
|
|
|
| def _null_emotion(model: DiT, c: Conditions) -> Conditions: |
| batch = c.emotion.shape[0] |
| return replace( |
| c, |
| emotion=torch.full( |
| (batch,), model.cfg.n_emotions, dtype=c.emotion.dtype, device=c.emotion.device |
| ), |
| intensity_known=torch.zeros(batch, dtype=torch.bool, device=c.emotion.device), |
| ) |
|
|
|
|
| def _concat(parts: list[Conditions]) -> Conditions: |
| """Concatenate along the batch; an optional field must be set in all parts or none.""" |
| values: dict[str, Tensor | None] = {} |
| for name in Conditions.__dataclass_fields__: |
| tensors = [getattr(p, name) for p in parts] |
| if all(x is None for x in tensors): |
| values[name] = None |
| elif any(x is None for x in tensors): |
| raise ValueError(f"Conditions field {name!r} is set in some parts only") |
| else: |
| values[name] = torch.cat(tensors) |
| return Conditions(**values) |
|
|
|
|
| def guided_velocity( |
| model: DiT, |
| x_t: Tensor, |
| t: Tensor, |
| c: Conditions, |
| *, |
| w_text: float, |
| w_emo: float, |
| drop_prompt_uncond: bool = False, |
| ) -> Tensor: |
| """Guided velocity ``(B, T, C)`` from one batched forward of ``3·B`` items.""" |
| batch = x_t.shape[0] |
| c_te = c |
| c_t = _null_emotion(model, c) |
| c_u = replace(c_t, drop_text=torch.ones(batch, dtype=torch.bool, device=x_t.device)) |
| if drop_prompt_uncond: |
| c_u = replace(c_u, cond_latent=torch.zeros_like(c.cond_latent)) |
| v = model(torch.cat([x_t] * 3), torch.cat([t] * 3), _concat([c_u, c_t, c_te])) |
| v_u, v_t, v_te = v.split(batch) |
| guided: Tensor = v_u + w_text * (v_t - v_u) + w_emo * (v_te - v_t) |
| return guided |
|
|
|
|
| @torch.no_grad() |
| def synthesize_latents( |
| model: DiT, |
| c: Conditions, |
| prompt_mask: Tensor, |
| *, |
| steps: int, |
| w_text: float, |
| w_emo: float, |
| seed: int, |
| method: Literal["euler", "midpoint"] = "euler", |
| sway: float = -1.0, |
| drop_prompt_uncond: bool = False, |
| ) -> Tensor: |
| """Generate latents ``(B, T, C)``. |
| |
| ``prompt_mask`` ``(B, T)`` bool, True = prompt frame: copied from ``cond_latent`` in the |
| output (they are context, not generated). The noise comes from a CPU generator seeded |
| with ``seed`` (callers always record it). |
| """ |
| batch, length = c.frame_mask.shape |
| if prompt_mask.shape != (batch, length) or prompt_mask.dtype != torch.bool: |
| raise ValueError(f"prompt_mask must be bool ({batch}, {length})") |
| device = c.cond_latent.device |
| generator = torch.Generator(device="cpu").manual_seed(seed) |
| x0 = torch.randn(batch, length, model.cfg.latent_dim, generator=generator) |
| x0 = x0.to(device=device, dtype=c.cond_latent.dtype) |
| model.eval() |
| out = sample_ode( |
| lambda x, t: guided_velocity( |
| model, x, t, c, w_text=w_text, w_emo=w_emo, drop_prompt_uncond=drop_prompt_uncond |
| ), |
| x0, |
| steps=steps, |
| method=method, |
| sway=sway, |
| ) |
| result: Tensor = torch.where(prompt_mask[..., None], c.cond_latent, out) |
| result = result * c.frame_mask[..., None] |
| return result |
|
|
