File size: 3,935 Bytes
c0bce59 | 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 | """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) # type: ignore[arg-type] # required fields are never None
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
|