wkea/blockdiffusion-api

1from __future__ import annotations
2
3from typing import Any, Dict, Optional
4
5import torch
6from torch import nn
7
8
9class CFGuidanceModel(nn.Module):
10    """
11    Classifier-Free Guidance 采样包装器：
12    eps = eps_uncond + scale * (eps_cond - eps_uncond)
13
14    说明：
15    - base_model 必须支持 forward(x, t, **model_kwargs)
16    - cond_kwargs/uncond_kwargs 用于传 tokens/mask 等条件信息
17    """
18
19    def __init__(
20        self,
21        base_model: nn.Module,
22        cond_kwargs: Optional[Dict[str, Any]] = None,
23        uncond_kwargs: Optional[Dict[str, Any]] = None,
24        scale: float = 3.0,
25    ):
26        super().__init__()
27        self.base_model = base_model
28        self.cond_kwargs = dict(cond_kwargs or {})
29        self.uncond_kwargs = dict(uncond_kwargs or {})
30        self.scale = float(scale)
31
32    def forward(self, x: torch.Tensor, t: torch.Tensor, **_: Any) -> torch.Tensor:
33        eps_uncond = self.base_model(x, t, **self.uncond_kwargs)
34        eps_cond = self.base_model(x, t, **self.cond_kwargs)
35        return eps_uncond + self.scale * (eps_cond - eps_uncond)
36
37
38