mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-10-12 02:17:16 +00:00
feat: Anima LLLite inpainting support (#2378)
* feat(inpainting): 4ch conditioning support with mask image handling and training adjustments * refactor(mask_generator): 政策の変更によりポリゴンと形状マスクの中心をキャンバス全体に均等に配置 refactor(inpainting): マスクチャンネルを[-1, 1]に正規化してRGB範囲に一致させる * fix(imports): accelerator_setupをインポートに追加 * fix: support for multi-GPU distributed inference with PartialState for Anima training * Add per-block torch.compile support for Anima ControlNet-LLLite Wire the shared --compile / --compile_* per-block torch.compile path (already used by anima_train_network.py) into the LLLite trainer, and fix two LLLite-specific patterns that defeated dynamo's per-block graph reuse and caused per-step recompilation: - Move the per-module depth-embedding addition out of the compiled forward into set_cond_image. The forward previously indexed depth_embeds[self.layer_idx] with a per-block Python int, which dynamo bakes as a static constant, forcing a distinct graph per block. The addition is now done outside the compiled region (still inside the autograd scope, so grad to depth_embeds is unchanged); numerically equivalent (depth_embeds is zero-init). Drops the now-dead _depth_embeds_ref attribute. - Disable force_nn_module_property_static_shapes in compile_transformer. LLLite injects the conditioning via the nn.Module attribute self.cond_emb, whose shape dynamo specializes statically by default even with dynamic=True (which only applies to forward inputs), causing a recompile per resolution bucket. Allowing module-property shapes to be dynamic collapses this to a single warmup compile. Also add the --compile x --torch_compile and --compile_fullgraph x --split_attn exclusivity asserts, apply CUDA perf switches, and document the feature in docs/anima_train_control_net_lllite.md (links to anima_torch_compile.md). Co-Authored-By: Claude Opus 4.8 <[email protected]> * Address Copilot review on Anima LLLite inpainting - LLLite: share cond embedding cx across modules and add per-module depth embedding inside forward, avoiding N copies of (B, S, cond_emb_dim). depth_embeds is re-indexed per step in set_cond_image (outside the compiled region) so multi-step backward keeps flowing grads and torch.compile still emits a single graph. - LLLite: emit the inpaint log suffix only for cond_in_channels == 4, consistent with the rest of the inpaint detection. - sampling: strip the --mk mask path like the --i image flag. - docs: fix the mask convention to state the mask is rescaled to {-1, 1} via mask * 2 - 1 (EN/JA), matching the implementation. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
Kohya S.
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
ac67c122a2
commit
ce2581e47a
@@ -74,6 +74,36 @@ def _load_control_image(
|
||||
return t.to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _load_mask_image(
|
||||
path: str, height: int, width: int, device: torch.device, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
"""Load and binarize a mask image to a (1, 1, H, W) tensor in {0, 1}.
|
||||
1.0 = inpaint area (穴), 0.0 = keep.
|
||||
"""
|
||||
img = Image.open(path).convert("L")
|
||||
if img.size != (width, height):
|
||||
img = img.resize((width, height), Image.NEAREST)
|
||||
arr = np.asarray(img).astype(np.float32) / 255.0
|
||||
arr = (arr >= 0.5).astype(np.float32)
|
||||
t = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).contiguous()
|
||||
return t.to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _build_inpaint_cond_image(
|
||||
rgb: torch.Tensor, mask: torch.Tensor, masked_input: bool
|
||||
) -> torch.Tensor:
|
||||
"""rgb: (B, 3, H, W) in [-1, 1], mask: (B, 1, H, W) in {0, 1}. Return (B, 4, H, W).
|
||||
|
||||
The mask channel is normalized to [-1, 1] (= (mask - 0.5) * 2) to match the RGB range.
|
||||
"""
|
||||
if masked_input:
|
||||
keep = (mask < 0.5).to(rgb.dtype)
|
||||
rgb = rgb * keep
|
||||
# mask channel: {0, 1} -> {-1, 1}. matches transforms.Normalize([0.5], [0.5])
|
||||
mask_pm1 = mask.to(rgb.dtype) * 2.0 - 1.0
|
||||
return torch.cat([rgb, mask_pm1], dim=1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# parse_args (replaces ami.parse_args)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -162,6 +192,21 @@ def parse_args() -> argparse.Namespace:
|
||||
"--lllite_use_aspp", type=str, default=None, choices=["true", "false"],
|
||||
help="override use_aspp from weights metadata (true/false)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mask_image", type=str, default=None,
|
||||
help=(
|
||||
"[inpaint] global mask image. Required for single-prompt inpainting (cond_in_channels=4). "
|
||||
"Per-prompt override: --mk <path>."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lllite_cond_in_channels", type=int, default=None,
|
||||
help="override cond_in_channels from weights metadata (3 or 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lllite_inpaint_masked_input", type=str, default=None, choices=["true", "false"],
|
||||
help="override inpaint_masked_input from weights metadata (true/false)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -222,6 +267,8 @@ def parse_prompt_line(line: str) -> Dict[str, Any]:
|
||||
overrides["negative_prompt"] = value
|
||||
elif option == "cn":
|
||||
overrides["control_image"] = value
|
||||
elif option == "mk":
|
||||
overrides["mask_image"] = value
|
||||
elif option == "am":
|
||||
overrides["lllite_multiplier"] = float(value)
|
||||
|
||||
@@ -275,12 +322,26 @@ def load_dit_model(args, device, dit_weight_dtype=None):
|
||||
else:
|
||||
from networks.control_net_lllite_anima import ASPP_DEFAULT_DILATIONS as _ASPP_DD
|
||||
aspp_dilations = _ASPP_DD
|
||||
cond_in_channels = (
|
||||
args.lllite_cond_in_channels
|
||||
if args.lllite_cond_in_channels is not None
|
||||
else int(meta.get("lllite.cond_in_channels", 3))
|
||||
)
|
||||
if args.lllite_inpaint_masked_input is not None:
|
||||
inpaint_masked_input = args.lllite_inpaint_masked_input == "true"
|
||||
else:
|
||||
inpaint_masked_input = (
|
||||
meta.get("lllite.inpaint_masked_input", "false").lower() == "true"
|
||||
)
|
||||
version = meta.get("lllite.version", "?")
|
||||
inpaint_log = (
|
||||
f", inpaint=on(masked_input={inpaint_masked_input})" if cond_in_channels == 4 else ""
|
||||
)
|
||||
logger.info(
|
||||
f"LLLite config (v{version}): cond_emb_dim={cond_emb_dim}, mlp_dim={mlp_dim}, "
|
||||
f"target_layers={target_layers}, cond_dim={cond_dim}, cond_resblocks={cond_resblocks}, "
|
||||
f"use_aspp={use_aspp}{(' dilations=' + str(list(aspp_dilations))) if use_aspp else ''}, "
|
||||
f"multiplier={args.lllite_multiplier}"
|
||||
f"cond_in_channels={cond_in_channels}{inpaint_log}, multiplier={args.lllite_multiplier}"
|
||||
)
|
||||
|
||||
lllite = ControlNetLLLiteDiT(
|
||||
@@ -293,6 +354,8 @@ def load_dit_model(args, device, dit_weight_dtype=None):
|
||||
cond_resblocks=cond_resblocks,
|
||||
use_aspp=use_aspp,
|
||||
aspp_dilations=aspp_dilations,
|
||||
cond_in_channels=cond_in_channels,
|
||||
inpaint_masked_input=inpaint_masked_input,
|
||||
)
|
||||
load_lllite_weights(lllite, args.lllite_weights, strict=False)
|
||||
lllite.apply_to()
|
||||
@@ -333,6 +396,23 @@ def generate_body(
|
||||
if not hasattr(anima, "lllite"):
|
||||
raise RuntimeError("DiT has no .lllite attribute; load_dit_model patch was not applied")
|
||||
|
||||
# inpainting (4ch): require a mask image; concat to cond_image as 4th channel
|
||||
if anima.lllite.cond_in_channels == 4:
|
||||
mk_path = getattr(args, "mask_image", None)
|
||||
if mk_path is None:
|
||||
raise ValueError(
|
||||
"mask_image is required for 4-channel (inpaint) LLLite. "
|
||||
"Specify --mask_image globally, or --mk per prompt in --from_file mode."
|
||||
)
|
||||
mask = _load_mask_image(mk_path, height, width, device, torch.bfloat16)
|
||||
cond_image = _build_inpaint_cond_image(
|
||||
cond_image, mask, anima.lllite.inpaint_masked_input
|
||||
)
|
||||
logger.info(
|
||||
f"Loaded mask image: {mk_path} -> 4ch cond_image {tuple(cond_image.shape)}"
|
||||
f" (masked_input={anima.lllite.inpaint_masked_input})"
|
||||
)
|
||||
|
||||
# honor per-prompt override of multiplier
|
||||
anima.lllite.set_multiplier(args.lllite_multiplier)
|
||||
anima.lllite.set_cond_image(cond_image)
|
||||
|
||||
@@ -37,6 +37,7 @@ from library import (
|
||||
)
|
||||
import library.accelerator_setup as accelerator_setup
|
||||
import library.args as args_util
|
||||
import library.compile_utils as compile_utils
|
||||
import library.dataset as dataset_util
|
||||
import library.model_io as model_io
|
||||
import library.optimizer as optimizer_util
|
||||
@@ -59,6 +60,7 @@ from networks.control_net_lllite_anima import (
|
||||
PRESETS as LLLITE_PRESETS,
|
||||
ATOMIC_SPECIFIERS as LLLITE_ATOMIC_SPECIFIERS,
|
||||
)
|
||||
from library.mask_generator import random_mask as _gen_random_mask
|
||||
|
||||
setup_logging()
|
||||
import logging
|
||||
@@ -74,12 +76,62 @@ def _load_control_image(path: str, width: int, height: int, device, dtype) -> to
|
||||
return tensor.to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _load_mask_image(path: str, width: int, height: int, device, dtype) -> torch.Tensor:
|
||||
"""Load a mask image and return (1, 1, H, W) in {0, 1}.
|
||||
1.0 = inpaint area (穴), 0.0 = keep.
|
||||
"""
|
||||
img = Image.open(path).convert("L").resize((width, height), Image.NEAREST)
|
||||
arr = np.array(img, dtype=np.float32) / 255.0
|
||||
arr = (arr >= 0.5).astype(np.float32)
|
||||
tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).contiguous() # (1, 1, H, W)
|
||||
return tensor.to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _build_inpaint_cond_image(
|
||||
rgb: torch.Tensor,
|
||||
masks: torch.Tensor,
|
||||
masked_input: bool,
|
||||
) -> torch.Tensor:
|
||||
"""rgb: (B, 3, H, W) in [-1, 1], masks: (B, 1, H, W) in {0, 1} (1=inpaint).
|
||||
Returns (B, 4, H, W) with the mask channel normalized to [-1, 1] to match the RGB range.
|
||||
|
||||
masked_input=True のとき、RGB を mask 域で 0 に潰してから concat する。
|
||||
"""
|
||||
if masked_input:
|
||||
keep = (masks < 0.5).to(rgb.dtype) # (B, 1, H, W)
|
||||
rgb = rgb * keep
|
||||
# mask channel: {0, 1} -> {-1, 1} (= (mask - 0.5) * 2). matches transforms.Normalize([0.5], [0.5])
|
||||
mask_pm1 = masks.to(rgb.dtype) * 2.0 - 1.0
|
||||
return torch.cat([rgb, mask_pm1], dim=1)
|
||||
|
||||
|
||||
def _generate_random_masks_for_batch(
|
||||
batch_size: int, height: int, width: int, device, dtype
|
||||
) -> torch.Tensor:
|
||||
"""library.mask_generator.random_mask を使ってバッチ分のランダム mask を生成する.
|
||||
返り値: (B, 1, H, W) in {0, 1} (1=inpaint, 0=keep)."""
|
||||
masks_np = np.empty((batch_size, 1, height, width), dtype=np.float32)
|
||||
for i in range(batch_size):
|
||||
pil = _gen_random_mask(width, height) # PIL "L", 0=keep / 255=inpaint
|
||||
arr = np.asarray(pil, dtype=np.float32) / 255.0
|
||||
arr = (arr >= 0.5).astype(np.float32)
|
||||
masks_np[i, 0] = arr
|
||||
return torch.from_numpy(masks_np).to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _make_lllite_sample_hooks(args, lllite, dit_dtype):
|
||||
"""Build (on_prompt_start, on_prompt_end) callbacks that wire control image / multiplier
|
||||
into the LLLite module before each sample prompt is rendered. The pre-sample multiplier is
|
||||
saved and restored so that, e.g., `--am 0` for inspection does not leak into training (which
|
||||
would otherwise hit the multiplier==0 short-circuit in LLLiteModuleDiT.forward and break
|
||||
backward by yielding a graph with no grad)."""
|
||||
backward by yielding a graph with no grad).
|
||||
|
||||
In inpainting mode (cond_in_channels=4) the prompt line additionally accepts `--mk <path>`
|
||||
for the mask image. If the mask is missing/not found the prompt is rendered without LLLite cond
|
||||
(warning logged), so users can intentionally inspect the base DiT.
|
||||
"""
|
||||
|
||||
is_inpaint = lllite.cond_in_channels == 4
|
||||
|
||||
saved = {"multiplier": None}
|
||||
|
||||
@@ -113,7 +165,28 @@ def _make_lllite_sample_hooks(args, lllite, dit_dtype):
|
||||
h = prompt_dict.get("height", 512)
|
||||
h = max(64, h - h % 16)
|
||||
w = max(64, w - w % 16)
|
||||
cond_image = _load_control_image(ci_path, w, h, accelerator.device, dit_dtype)
|
||||
rgb = _load_control_image(ci_path, w, h, accelerator.device, dit_dtype)
|
||||
|
||||
if is_inpaint:
|
||||
mk_path = prompt_dict.get("mask_image")
|
||||
if mk_path is None:
|
||||
logger.warning(
|
||||
"inpaint LLLite: no mask image for sample prompt (use '--mk <path>'); "
|
||||
"running base DiT without LLLite cond"
|
||||
)
|
||||
lllite.clear_cond_image()
|
||||
return
|
||||
if not os.path.isfile(mk_path):
|
||||
logger.warning(
|
||||
f"inpaint LLLite: mask image not found: {mk_path}; running base DiT without LLLite cond"
|
||||
)
|
||||
lllite.clear_cond_image()
|
||||
return
|
||||
mask = _load_mask_image(mk_path, w, h, accelerator.device, dit_dtype)
|
||||
cond_image = _build_inpaint_cond_image(rgb, mask, args.lllite_inpaint_masked_input)
|
||||
else:
|
||||
cond_image = rgb
|
||||
|
||||
lllite.set_cond_image(cond_image)
|
||||
|
||||
def on_prompt_end(prompt_dict: dict):
|
||||
@@ -184,6 +257,25 @@ def add_anima_lllite_arguments(parser: argparse.ArgumentParser):
|
||||
default=None,
|
||||
help="pretrained LLLite weights to resume from / 学習を再開する LLLite の初期重み",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lllite_cond_in_channels",
|
||||
type=int,
|
||||
default=3,
|
||||
help=(
|
||||
"number of input channels for the LLLite conditioning1 trunk (default: 3, RGB only). "
|
||||
"Set to 4 to enable inpainting mode (RGB + 1ch mask). "
|
||||
"/ LLLite の conditioning1 入力チャネル数。デフォルト 3 (RGB)、4 で inpainting 用 (RGB+mask)"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lllite_inpaint_masked_input",
|
||||
action="store_true",
|
||||
help=(
|
||||
"[inpaint] additionally zero out RGB inside the mask region before concatenating with mask. "
|
||||
"Only effective when --lllite_cond_in_channels=4. "
|
||||
"/ inpainting 時、RGB の mask 域を 0 で穴埋めしてから concat する (cond_in_channels=4 のときのみ有効)"
|
||||
),
|
||||
)
|
||||
# --conditioning_data_dir は args_util.add_dataset_arguments 側で既に定義済み
|
||||
|
||||
|
||||
@@ -217,6 +309,17 @@ def train(args):
|
||||
"fused_backward_pass is not supported in Anima ControlNet-LLLite training (MVP)"
|
||||
)
|
||||
|
||||
# per-block torch.compile の排他チェック (anima_train_network.assert_extra_args と同等)
|
||||
if args.compile:
|
||||
assert not args.torch_compile, (
|
||||
"--compile (per-block torch.compile) and --torch_compile (accelerate dynamo) cannot be used together"
|
||||
" / --compile（ブロック単位torch.compile）と--torch_compile（accelerate dynamo）は併用できません"
|
||||
)
|
||||
assert not (args.compile_fullgraph and args.split_attn), (
|
||||
"--compile_fullgraph cannot be used with --split_attn (split attention uses dynamic control flow)"
|
||||
" / --compile_fullgraphは--split_attnと併用できません（split attentionは動的な制御フローを使用します）"
|
||||
)
|
||||
|
||||
cache_latents = args.cache_latents
|
||||
|
||||
if args.seed is not None:
|
||||
@@ -382,6 +485,16 @@ def train(args):
|
||||
|
||||
dit.requires_grad_(False)
|
||||
|
||||
# inpainting (4ch) フラグの早期検証
|
||||
if args.lllite_cond_in_channels < 1:
|
||||
raise ValueError(f"--lllite_cond_in_channels must be >= 1, got {args.lllite_cond_in_channels}")
|
||||
if args.lllite_inpaint_masked_input and args.lllite_cond_in_channels != 4:
|
||||
logger.warning(
|
||||
f"--lllite_inpaint_masked_input is only effective when --lllite_cond_in_channels=4 "
|
||||
f"(got {args.lllite_cond_in_channels}); flag will be ignored at runtime"
|
||||
)
|
||||
is_inpaint = args.lllite_cond_in_channels == 4
|
||||
|
||||
# Build LLLite (DiT を走査して各 Attention Linear に貼る)
|
||||
logger.info("Building ControlNet-LLLite (Anima)...")
|
||||
lllite = ControlNetLLLiteDiT(
|
||||
@@ -394,6 +507,8 @@ def train(args):
|
||||
cond_dim=args.lllite_cond_dim,
|
||||
cond_resblocks=args.lllite_cond_resblocks,
|
||||
use_aspp=args.lllite_use_aspp,
|
||||
cond_in_channels=args.lllite_cond_in_channels,
|
||||
inpaint_masked_input=args.lllite_inpaint_masked_input,
|
||||
)
|
||||
|
||||
if args.network_weights is not None:
|
||||
@@ -473,6 +588,17 @@ def train(args):
|
||||
if args.full_fp16:
|
||||
accelerator_setup.patch_accelerator_for_fp16_training(accelerator)
|
||||
|
||||
# CUDA perf switches are independent of torch.compile; apply whenever requested.
|
||||
compile_utils.apply_cuda_optimizations(args)
|
||||
|
||||
if args.compile:
|
||||
# per-block torch.compile を凍結 DiT のブロックに適用する。LLLite の forward 差し替え
|
||||
# (apply_to) と accelerator.prepare の後でなければならない。block swap は MVP で無効
|
||||
# のため disable_linear=False 固定。LLLite モジュールは対象 Linear の forward を差し替え
|
||||
# ているため、compile は patch 済みの forward を取り込む (cond_emb はガード付き入力扱い)。
|
||||
dit_to_compile = accelerator.unwrap_model(wrapper).dit
|
||||
compile_utils.compile_transformer(args, dit_to_compile, [dit_to_compile.blocks], disable_linear=False)
|
||||
|
||||
args_util.resume_from_local_or_hf_if_specified(accelerator, args)
|
||||
|
||||
# epoch計算
|
||||
@@ -553,6 +679,10 @@ def train(args):
|
||||
sai_metadata["lllite.use_aspp"] = "true" if args.lllite_use_aspp else "false"
|
||||
if args.lllite_use_aspp:
|
||||
sai_metadata["lllite.aspp_dilations"] = ",".join(str(d) for d in unwrapped.aspp_dilations)
|
||||
sai_metadata["lllite.cond_in_channels"] = str(args.lllite_cond_in_channels)
|
||||
sai_metadata["lllite.inpaint_masked_input"] = (
|
||||
"true" if args.lllite_inpaint_masked_input else "false"
|
||||
)
|
||||
save_lllite_model(ckpt_file, unwrapped, dtype=save_dtype, metadata=sai_metadata)
|
||||
|
||||
def _save_step(global_step_: int, epoch_: int):
|
||||
@@ -658,6 +788,16 @@ def train(args):
|
||||
# cond image: dataset 側で IMAGE_TRANSFORMS により [-1,1] 正規化済み
|
||||
cond_image = batch["conditioning_images"].to(accelerator.device, dtype=dit_weight_dtype)
|
||||
|
||||
# inpainting: ランダム mask をバッチ毎に生成し、cond_image を 4ch (RGB + mask) 化
|
||||
if is_inpaint:
|
||||
bs_c, _, h_c, w_c = cond_image.shape
|
||||
mask = _generate_random_masks_for_batch(
|
||||
bs_c, h_c, w_c, accelerator.device, dit_weight_dtype
|
||||
)
|
||||
cond_image = _build_inpaint_cond_image(
|
||||
cond_image, mask, args.lllite_inpaint_masked_input
|
||||
)
|
||||
|
||||
# 5D化
|
||||
noisy_model_input = noisy_model_input.unsqueeze(2) # (B, C, 1, H, W)
|
||||
|
||||
|
||||
@@ -352,6 +352,8 @@ The metadata records `modelspec.architecture = "anima-preview/control-net-lllite
|
||||
| `lllite.mlp_dim` | per-module MLP / FiLM hidden dim |
|
||||
| `lllite.target_layers` | the user-supplied `--lllite_target_layers` string verbatim |
|
||||
| `lllite.target_atomics` | resolved canonical atomic specifier list (comma-separated) |
|
||||
| `lllite.cond_in_channels` | conditioning1 input channel count (`"3"` standard, `"4"` inpaint) |
|
||||
| `lllite.inpaint_masked_input` | `"true"` if RGB was masked before concat at train time (inpaint only) |
|
||||
|
||||
The inference script (Section 6) reads these back, so you normally do not need to specify the architecture options on the command line again. Currently the inference script **requires** the metadata to reconstruct the architecture: state-dict-only auto-detection (e.g. for hand-edited weight files) is not implemented.
|
||||
|
||||
@@ -402,6 +404,8 @@ lllite_conditioning1.out_norm.{weight,bias}
|
||||
| `lllite.mlp_dim` | LLLite モジュール毎の MLP / FiLM 中間次元 |
|
||||
| `lllite.target_layers` | ユーザが指定した `--lllite_target_layers` 文字列をそのまま記録 |
|
||||
| `lllite.target_atomics` | 解決後の canonical atomic specifier 列（カンマ区切り） |
|
||||
| `lllite.cond_in_channels` | conditioning1 入力チャネル数（`"3"` 通常、`"4"` inpainting） |
|
||||
| `lllite.inpaint_masked_input` | 学習時 RGB を mask 域で 0 化してから concat したか（inpaint 専用、`"true"`/`"false"`） |
|
||||
|
||||
これらは推論スクリプト（第 6 節）で自動的に読み出されるため、通常はコマンドラインで再指定する必要はありません。現在の推論スクリプトはアーキテクチャ復元に**メタデータが必須**で、state_dict 単独からの自動判定（手編集された重みファイル等）には対応していません。
|
||||
|
||||
@@ -480,12 +484,144 @@ CFG 推論（cond / uncond の 2 pass）は両 pass に同じ `cond_emb` を配
|
||||
|
||||
</details>
|
||||
|
||||
## 7. Tips & Limitations / 補足と制限
|
||||
## 7. Inpainting Support / Inpainting 対応
|
||||
|
||||
> **Status:** experimental. Adds an optional 4-channel conditioning path (RGB + 1ch binary mask). When enabled, LLLite is trained to fill in the masked region of an image given the rest of it as conditioning.
|
||||
|
||||
When `--lllite_cond_in_channels=4` is passed, the LLLite `conditioning1` trunk accepts a 4-channel input instead of the default 3 channels: `[R, G, B, mask]`. The training script generates a fresh random mask per sample at every step (via `library.mask_generator.random_mask`) and concatenates it with the conditioning RGB image. At inference time you pass the control image and the mask separately (the script concatenates them internally).
|
||||
|
||||
The default 3-channel behavior (no mask, generic ControlNet usage) is unchanged. The same code path is intentionally generic so that future 4-channel conditioning (e.g. user-supplied control masks for lighting / region control) can reuse the same trunk; for the initial release only inpainting is wired up.
|
||||
|
||||
### 7.1. Mask Convention / マスク規約
|
||||
|
||||
* `white` (1.0) = **inpaint area** — the region the model should fill in / 生成すべき穴.
|
||||
* `black` (0.0) = **keep** — the region whose pixels are preserved.
|
||||
|
||||
The mask is binarized at 0.5 on load (`{0.0, 1.0}`) and then rescaled to `{-1.0, 1.0}` via `mask * 2 - 1` before being concatenated as the 4th channel, so that the mask channel matches the `[-1, 1]` range of the RGB conditioning channels.
|
||||
|
||||
### 7.2. Dataset Setup / データセット設定
|
||||
|
||||
For inpainting, **the conditioning image is the original training image itself**. The simplest way to wire this up without modifying the dataset code is to point `conditioning_data_dir` at the same directory as `image_dir` (or a symlink / copy of it). The same image is then loaded twice per step (once as the training target, once as the LLLite conditioning), and the random mask is generated on the fly inside the training script.
|
||||
|
||||
Minimal TOML:
|
||||
|
||||
```toml
|
||||
[general]
|
||||
caption_extension = ".txt"
|
||||
shuffle_caption = false
|
||||
|
||||
[[datasets]]
|
||||
resolution = 1024
|
||||
batch_size = 1
|
||||
|
||||
[[datasets.subsets]]
|
||||
image_dir = "/path/to/training_images"
|
||||
conditioning_data_dir = "/path/to/training_images" # same as image_dir for inpainting
|
||||
num_repeats = 1
|
||||
```
|
||||
|
||||
The CLI form `--train_data_dir <dir> --conditioning_data_dir <dir>` works analogously (point both at the same directory).
|
||||
|
||||
### 7.3. Training-Specific Arguments / 学習引数
|
||||
|
||||
* `--lllite_cond_in_channels=<int>` (default `3`)
|
||||
* `3` = standard ControlNet-LLLite (RGB-only conditioning, no mask).
|
||||
* `4` = inpainting (RGB + 1ch mask). Activates random-mask generation per step.
|
||||
|
||||
* `--lllite_inpaint_masked_input` (flag, default off)
|
||||
* Only effective when `--lllite_cond_in_channels=4`. If set, the RGB channels of the conditioning image are zeroed in the mask region **before** being concatenated with the mask channel (so the model never sees the ground-truth pixels behind the hole). Without this flag, the model sees the full RGB image plus the mask channel — the model is expected to learn to ignore the masked region from the mask signal. The flag is recorded in the saved metadata as `lllite.inpaint_masked_input`.
|
||||
|
||||
Example command (one line in practice; line continuations as in Section 3):
|
||||
|
||||
```bash
|
||||
accelerate launch --num_cpu_threads_per_process 1 anima_train_control_net_lllite.py \
|
||||
--pretrained_model_name_or_path="<path to Anima DiT model>" \
|
||||
--qwen3="<path to Qwen3-0.6B model or directory>" \
|
||||
--vae="<path to Qwen-Image VAE model>" \
|
||||
--dataset_config="my_anima_lllite_inpaint.toml" \
|
||||
--output_dir="<output directory>" \
|
||||
--output_name="my_anima_lllite_inpaint" \
|
||||
--save_model_as=safetensors \
|
||||
--cond_emb_dim=32 --lllite_mlp_dim=64 --lllite_cond_dim=64 \
|
||||
--lllite_target_layers=self_attn_q \
|
||||
--lllite_cond_in_channels=4 \
|
||||
--learning_rate=1e-4 --optimizer_type="AdamW8bit" --lr_scheduler="constant" \
|
||||
--timestep_sampling="shift" --discrete_flow_shift=3.0 \
|
||||
--max_train_epochs=10 --save_every_n_epochs=1 \
|
||||
--mixed_precision="bf16" --gradient_checkpointing \
|
||||
--cache_latents --cache_text_encoder_outputs \
|
||||
--vae_chunk_size=64 --vae_disable_cache
|
||||
```
|
||||
|
||||
### 7.4. Sample Image Generation / 学習中のサンプル生成
|
||||
|
||||
The sample prompt line accepts an additional flag in inpainting mode:
|
||||
|
||||
* `--mk <path>` — per-prompt mask image (white=inpaint, black=keep).
|
||||
|
||||
In inpaint mode (`cond_in_channels=4`), if `--mk` is omitted (or the file is not found) the prompt is rendered with the base DiT (LLLite cond cleared) and a warning is logged. The `--cn <path>` flag is still required (it supplies the RGB conditioning image).
|
||||
|
||||
Example `prompts.txt`:
|
||||
|
||||
```
|
||||
a cat in a forest --w 1024 --h 1024 --cn original_a.png --mk mask_a.png --d 42
|
||||
a sunset over mountains --w 1024 --h 1024 --cn original_b.png --mk mask_b.png
|
||||
```
|
||||
|
||||
### 7.5. Inference / 推論
|
||||
|
||||
The inference script auto-detects the 4-channel mode from the saved metadata (`lllite.cond_in_channels`, `lllite.inpaint_masked_input`) and requires a mask alongside the control image:
|
||||
|
||||
* `--mask_image <path>` — global mask image (required for single-prompt inference when `cond_in_channels=4`).
|
||||
* Prompt-line `--mk <path>` — per-prompt mask override in `--from_file` / `--interactive` mode.
|
||||
* `--lllite_cond_in_channels {3,4}` / `--lllite_inpaint_masked_input {true,false}` — manual metadata overrides (usually unnecessary).
|
||||
|
||||
Single-prompt example:
|
||||
|
||||
```bash
|
||||
python anima_minimal_inference_control_net_lllite.py \
|
||||
--dit "<path to Anima DiT>" --vae "<path to Qwen-Image VAE>" --text_encoder "<path to Qwen3-0.6B>" \
|
||||
--lllite_weights "out/my_anima_lllite_inpaint-last.safetensors" \
|
||||
--control_image "original.png" --mask_image "mask.png" \
|
||||
--prompt "a cat in a forest" --image_size 1024 1024 \
|
||||
--infer_steps 50 --guidance_scale 3.5 --save_path "out/"
|
||||
```
|
||||
|
||||
### 7.6. Limitations / 制限
|
||||
|
||||
* The model architecture is otherwise unchanged — the DiT itself is not modified for inpainting. All "where to fill" information must flow through the LLLite conditioning path.
|
||||
* Only random masks are supported during training in this release; future versions may allow user-provided masks via dataset configuration.
|
||||
* The 4-channel feature is implementation-orthogonal to ASPP, ResBlock count, target_layers, etc. Combine freely.
|
||||
* `cond_in_channels` other than `3` and `4` are accepted by the architecture (it is generic), but only `4` is wired up to the random-mask training path. Custom values are intended for future user-supplied channels.
|
||||
|
||||
<details>
|
||||
<summary>日本語</summary>
|
||||
|
||||
`--lllite_cond_in_channels=4` を指定すると、`conditioning1` の入力チャネルが 3 → 4 になり、`[R, G, B, mask]` を受け付けるようになります。学習時は `library.mask_generator.random_mask` でステップ毎にランダム mask を生成し、`conditioning_images` (= 元画像) と concat して渡します。3ch のデフォルト動作は変更されません。
|
||||
|
||||
**マスク規約**: 白 (1.0) = inpaint 対象 (穴) / 黒 (0.0) = 保持。ロード時に 0.5 で二値化 (`{0.0, 1.0}`) した後、`mask * 2 - 1` で `{-1.0, 1.0}` にスケールしてから 4ch 目として concat します (RGB conditioning チャンネルの `[-1, 1]` レンジに合わせるため)。
|
||||
|
||||
**データセット**: inpainting では「conditioning 画像 = 元画像」となるため、`conditioning_data_dir` には `image_dir` と同じパス (もしくは symlink / コピー) を指定してください。データセットコード側は無改修です。
|
||||
|
||||
**学習引数**:
|
||||
- `--lllite_cond_in_channels=4`: inpainting モードを有効化 (ステップ毎ランダム mask 生成 + 4ch 化)。
|
||||
- `--lllite_inpaint_masked_input`: フラグ。立てると RGB を mask 域で 0 化してから concat します (穴の中の真値は見せない)。OFF だと RGB はそのまま渡し、mask チャネル経由で「どこが穴か」を学習させます。メタデータ `lllite.inpaint_masked_input` に記録されます。
|
||||
|
||||
**サンプル画像生成**: prompt 行に `--mk <path>` を追加できます。inpaint モードで `--mk` が指定されない／ファイルが見つからない場合は LLLite cond をクリアし素の DiT で生成し warning を出します。
|
||||
|
||||
**推論**: メタデータから自動的に 4ch モードを検出します。`--mask_image <path>` (グローバル) または prompt 行 `--mk <path>` (per-prompt) で mask を指定してください。
|
||||
|
||||
**制限**: DiT 本体は無改修なので、「どこを埋めるか」の情報は LLLite cond 経路を通る必要があります。現状ランダム mask のみ対応 (将来的にはユーザー指定 mask に対応予定)。`cond_in_channels` の `3`/`4` 以外も構築可能ですが、本リリースでは `4` のみが random-mask 学習パスに接続されています。
|
||||
|
||||
</details>
|
||||
|
||||
## 8. Tips & Limitations / 補足と制限
|
||||
|
||||
* **Resolution alignment.** The conditioning encoder uses fixed stride 16, so `cond_image` HW must equal `latent HW × 8` (i.e. the original training image size) (the latent is patchified with patch size=2, so stride is 8*2=16). The DataLoader for the ControlNet dataset already resizes the conditioning image to match the training image, so in practice you only need to make sure the control image you pass at inference time matches the requested `--image_size`.
|
||||
* **`T=1` only.** Video-style multi-frame inputs are not supported — the wrapper asserts `T==1` at forward time.
|
||||
* **Bucket size.** The training script enforces a bucket resolution step of 16 (Qwen-Image VAE /8 × patch /2).
|
||||
* **Memory.** `--blocks_to_swap`, `--cpu_offload_checkpointing`, `--unsloth_offload_checkpointing` are not yet supported. If VRAM is tight, prefer `--full_bf16`, smaller `--lllite_mlp_dim`, lower `--cond_emb_dim`, and `--gradient_checkpointing`.
|
||||
* **`torch.compile` (speed-up).** Per-block `torch.compile` of the DiT is supported here as well, via the same `--compile` / `--compile_*` arguments shared with the LoRA script. See [`anima_torch_compile.md`](./anima_torch_compile.md) for details and recommended settings. Note that with aspect-ratio bucketing each distinct bucket resolution triggers a one-time recompile, so the first pass over your buckets is slower before it settles.
|
||||
* **Save format.** The saved `.safetensors` is **not** compatible with the SDXL LLLite format and **not** loadable by `sdxl_gen_img.py`. Use the dedicated inference script in Section 6.
|
||||
* **Metadata-required at inference.** Inference relies on the architecture metadata (`lllite.version`, `lllite.cond_dim`, `lllite.cond_resblocks`, `lllite.use_aspp`, `lllite.target_atomics`, ...) saved by the training script to reconstruct the LLLite architecture. State-dict-only auto-detection of those fields is not implemented; if a weight file lacks metadata, you currently need to pass the override flags listed in Section 6.3 explicitly.
|
||||
|
||||
@@ -496,6 +632,7 @@ CFG 推論（cond / uncond の 2 pass）は両 pass に同じ `cond_emb` を配
|
||||
* **`T=1` のみ.** 動画的な多フレーム入力はサポートしていません（wrapper の forward 冒頭で assert）。
|
||||
* **bucket サイズ.** 学習スクリプトは bucket 解像度ステップを 16（Qwen-Image VAE /8 × patch /2）として検証します。
|
||||
* **メモリ.** `--blocks_to_swap`、`--cpu_offload_checkpointing`、`--unsloth_offload_checkpointing` は未対応です。VRAM が厳しい場合は `--full_bf16`、`--lllite_mlp_dim` を下げる、`--cond_emb_dim` を下げる、`--gradient_checkpointing` を有効にする、などで対応してください。
|
||||
* **`torch.compile`（高速化）.** LoRA 学習スクリプトと共通の `--compile` / `--compile_*` 引数により、本スクリプトでも DiT のブロック単位 `torch.compile` が利用できます。詳細や推奨設定は [`anima_torch_compile.md`](./anima_torch_compile.md) を参照してください。なお aspect-ratio bucketing 使用時はバケット解像度ごとに初回 1 回の recompile が走るため、全バケットを一巡するまでは遅くなりますが、その後は安定します。
|
||||
* **保存形式.** 保存される `.safetensors` は SDXL LLLite フォーマットとは**互換性がなく**、`sdxl_gen_img.py` ではロードできません。推論には第 6 節の専用スクリプトを使用してください。
|
||||
* **推論時のメタデータ依存.** 推論時のアーキテクチャ復元は、学習スクリプトが書き込んだメタデータ（`lllite.version` / `lllite.cond_dim` / `lllite.cond_resblocks` / `lllite.use_aspp` / `lllite.target_atomics` など）に依存します。state_dict 単独からの自動判定は実装されていないため、メタデータの無い重みを使う場合は第 6.3 節の手動上書き引数を明示的に指定してください。
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ from typing import Optional
|
||||
from library import device_utils
|
||||
import numpy as np
|
||||
import torch
|
||||
from accelerate import Accelerator
|
||||
from accelerate import Accelerator, PartialState
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
|
||||
@@ -549,6 +549,8 @@ def sample_images(
|
||||
|
||||
dit.switch_block_swap_for_inference()
|
||||
|
||||
distributed_state = PartialState() # for multi gpu distributed inference. this is a singleton, so it's safe to use it here
|
||||
|
||||
prompts = sampling.load_prompts(args.sample_prompts)
|
||||
save_dir = os.path.join(args.output_dir, "sample")
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
@@ -561,30 +563,44 @@ def sample_images(
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with torch.no_grad(), accelerator.autocast():
|
||||
for prompt_dict in prompts:
|
||||
dit.prepare_block_swap_before_forward()
|
||||
if on_prompt_start is not None:
|
||||
on_prompt_start(prompt_dict, accelerator)
|
||||
try:
|
||||
_sample_image_inference(
|
||||
accelerator,
|
||||
args,
|
||||
dit,
|
||||
text_encoder,
|
||||
vae,
|
||||
tokenize_strategy,
|
||||
text_encoding_strategy,
|
||||
save_dir,
|
||||
prompt_dict,
|
||||
epoch,
|
||||
steps,
|
||||
sample_prompts_te_outputs,
|
||||
prompt_replacement,
|
||||
)
|
||||
finally:
|
||||
if on_prompt_end is not None:
|
||||
on_prompt_end(prompt_dict)
|
||||
def _run_one(prompt_dict):
|
||||
dit.prepare_block_swap_before_forward()
|
||||
if on_prompt_start is not None:
|
||||
on_prompt_start(prompt_dict, accelerator)
|
||||
try:
|
||||
_sample_image_inference(
|
||||
accelerator,
|
||||
args,
|
||||
dit,
|
||||
text_encoder,
|
||||
vae,
|
||||
tokenize_strategy,
|
||||
text_encoding_strategy,
|
||||
save_dir,
|
||||
prompt_dict,
|
||||
epoch,
|
||||
steps,
|
||||
sample_prompts_te_outputs,
|
||||
prompt_replacement,
|
||||
)
|
||||
finally:
|
||||
if on_prompt_end is not None:
|
||||
on_prompt_end(prompt_dict)
|
||||
|
||||
if distributed_state.num_processes <= 1:
|
||||
with torch.no_grad(), accelerator.autocast():
|
||||
for prompt_dict in prompts:
|
||||
_run_one(prompt_dict)
|
||||
else:
|
||||
# Distribute prompts across processes. Each process handles prompts[i :: num_processes].
|
||||
per_process_prompts = [] # list of lists
|
||||
for i in range(distributed_state.num_processes):
|
||||
per_process_prompts.append(prompts[i :: distributed_state.num_processes])
|
||||
|
||||
with torch.no_grad(), accelerator.autocast():
|
||||
with distributed_state.split_between_processes(per_process_prompts) as prompt_dict_lists:
|
||||
for prompt_dict in prompt_dict_lists[0]:
|
||||
_run_one(prompt_dict)
|
||||
|
||||
# Restore RNG state
|
||||
torch.set_rng_state(rng_state)
|
||||
|
||||
@@ -86,6 +86,14 @@ def compile_transformer(
|
||||
if args.compile_cache_size_limit is not None:
|
||||
torch._dynamo.config.cache_size_limit = args.compile_cache_size_limit
|
||||
|
||||
# nn.Module の tensor 属性 (例: ControlNet-LLLite が注入する self.cond_emb) は、
|
||||
# 既定では dynamic=True でも shape が static に specialize され、解像度バケット毎に
|
||||
# recompile を誘発する。forward 入力ではなく属性経由で渡る可変長テンソルを dynamic
|
||||
# 対象に含めるため、この強制 static 化を無効化する。Parameter (学習重み) は形状不変
|
||||
# なので force_parameter_static_shapes は既定 (True) のままにしておく。
|
||||
if hasattr(torch._dynamo.config, "force_nn_module_property_static_shapes"):
|
||||
torch._dynamo.config.force_nn_module_property_static_shapes = False
|
||||
|
||||
for blocks in target_blocks:
|
||||
for i, block in enumerate(blocks):
|
||||
blocks[i] = torch.compile(
|
||||
|
||||
@@ -136,8 +136,11 @@ def polygon_mask(
|
||||
|
||||
for _ in range(n_polygons):
|
||||
radius = rng.uniform(min_r, max_r)
|
||||
cx = rng.uniform(radius * 0.2, width - radius * 0.2)
|
||||
cy = rng.uniform(radius * 0.2, height - radius * 0.2)
|
||||
# Center uniformly anywhere on the canvas; vertices that fall outside are clipped by PIL.
|
||||
# This is intentional: large polygons centered near edges/corners give edge-region masks,
|
||||
# which the previous radius*0.2 margin suppressed.
|
||||
cx = rng.uniform(0, width)
|
||||
cy = rng.uniform(0, height)
|
||||
pts = _random_convex_polygon(cx, cy, radius, n_points, irregularity, rng)
|
||||
draw.polygon(pts, fill=255)
|
||||
|
||||
@@ -171,8 +174,10 @@ def shape_mask(
|
||||
size_w = rng.randint(int(width * min_coverage ** 0.5), int(width * max_coverage ** 0.5))
|
||||
size_h = rng.randint(int(height * min_coverage ** 0.5), int(height * max_coverage ** 0.5))
|
||||
|
||||
x1 = rng.randint(0, width - size_w)
|
||||
y1 = rng.randint(0, height - size_h)
|
||||
# Allow the shape's bbox to overhang the canvas by up to half its size, so its center is
|
||||
# uniform in [0, width) x [0, height). Off-canvas portions are clipped by PIL.
|
||||
x1 = rng.randint(-size_w // 2, width - size_w // 2)
|
||||
y1 = rng.randint(-size_h // 2, height - size_h // 2)
|
||||
x2, y2 = x1 + size_w, y1 + size_h
|
||||
|
||||
use_ellipse = (shape == "ellipse") or (shape == "random" and rng.random() < 0.5)
|
||||
|
||||
+6
-1
@@ -8,7 +8,7 @@ call periodically (every N steps / epochs) to produce preview images:
|
||||
``--sample_sampler`` (ddim / pndm / euler / dpmsolver / ...).
|
||||
- :func:`line_to_prompt_dict` / :func:`load_prompts` — parse prompt files
|
||||
(``.txt`` / ``.toml`` / ``.json``) into prompt dicts. ``--w / --h /
|
||||
--d / --s / --l / --g / --n / --ss / --cn / --i / --ctr / --rcfg /
|
||||
--d / --s / --l / --g / --n / --ss / --cn / --mk / --i / --ctr / --rcfg /
|
||||
--fs / --am`` flags are recognised.
|
||||
- :func:`sample_images` / :func:`sample_images_common` /
|
||||
:func:`sample_image_inference` — orchestrate one round of sample
|
||||
@@ -178,6 +178,11 @@ def line_to_prompt_dict(line: str) -> dict:
|
||||
prompt_dict["controlnet_image"] = m.group(1)
|
||||
continue
|
||||
|
||||
m = re.match(r"mk (.+)", parg, re.IGNORECASE)
|
||||
if m:
|
||||
prompt_dict["mask_image"] = m.group(1).strip()
|
||||
continue
|
||||
|
||||
m = re.match(r"i (.+)", parg, re.IGNORECASE)
|
||||
if m:
|
||||
prompt_dict["image"] = m.group(1).strip()
|
||||
|
||||
@@ -144,7 +144,7 @@ class _ASPP(nn.Module):
|
||||
class _Conditioning1(nn.Module):
|
||||
"""v2 conditioning trunk.
|
||||
|
||||
in (B,3,H,W)
|
||||
in (B,C_in,H,W)
|
||||
-> Conv 4x4 s=4 + GN + SiLU # cond_dim/2, H/4
|
||||
-> Conv 3x3 s=1 + GN + SiLU # cond_dim/2, H/4 (受容野拡張)
|
||||
-> Conv 4x4 s=4 + GN + SiLU # cond_dim, H/16 (token 解像度)
|
||||
@@ -152,6 +152,8 @@ class _Conditioning1(nn.Module):
|
||||
-> Conv 1x1 # cond_emb_dim
|
||||
-> flatten (B, S, cond_emb_dim)
|
||||
-> LayerNorm
|
||||
|
||||
C_in は cond_in_channels で指定する。デフォルト 3 (RGB のみ)、4 で inpainting (RGB+mask) 等。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -161,12 +163,15 @@ class _Conditioning1(nn.Module):
|
||||
n_resblocks: int,
|
||||
use_aspp: bool = False,
|
||||
aspp_dilations: Tuple[int, ...] = ASPP_DEFAULT_DILATIONS,
|
||||
cond_in_channels: int = 3,
|
||||
):
|
||||
super().__init__()
|
||||
assert cond_dim % 2 == 0, f"cond_dim must be even, got {cond_dim}"
|
||||
assert cond_in_channels >= 1, f"cond_in_channels must be >= 1, got {cond_in_channels}"
|
||||
ch_half = cond_dim // 2
|
||||
|
||||
self.conv1 = nn.Conv2d(3, ch_half, kernel_size=4, stride=4, padding=0)
|
||||
self.cond_in_channels = cond_in_channels
|
||||
self.conv1 = nn.Conv2d(cond_in_channels, ch_half, kernel_size=4, stride=4, padding=0)
|
||||
self.norm1 = _gn(ch_half)
|
||||
self.conv2 = nn.Conv2d(ch_half, ch_half, kernel_size=3, stride=1, padding=1)
|
||||
self.norm2 = _gn(ch_half)
|
||||
@@ -234,12 +239,16 @@ class LLLiteModuleDiT(nn.Module):
|
||||
nn.init.zeros_(self.up.weight)
|
||||
nn.init.zeros_(self.up.bias)
|
||||
|
||||
# 親 ControlNetLLLiteDiT が set_cond_image で注入する
|
||||
# 親 ControlNetLLLiteDiT が set_cond_image で注入する。
|
||||
# cond_emb は全モジュールで共有される cx (B, S, cond_emb_dim)、
|
||||
# depth_emb はこのモジュール用の depth embedding (cond_emb_dim,)。
|
||||
# cx を共有参照することで N コピーを避け、加算は forward 内で行う。
|
||||
self.cond_emb: Optional[torch.Tensor] = None
|
||||
self.depth_emb: Optional[torch.Tensor] = None
|
||||
|
||||
# 親 ControlNetLLLiteDiT が __init__ 末尾で設定 (list 包みで Parameter 二重登録回避)
|
||||
# 親 ControlNetLLLiteDiT が __init__ 末尾で layer_idx を設定する。
|
||||
# depth embedding の index 参照は set_cond_image 側で行う (torch.compile 対策)。
|
||||
self.layer_idx: int = -1
|
||||
self._depth_embeds_ref: List[nn.Parameter] = []
|
||||
|
||||
def apply_to(self):
|
||||
self.org_forward = self.org_module[0].forward
|
||||
@@ -259,24 +268,21 @@ class LLLiteModuleDiT(nn.Module):
|
||||
B, T, H, W, D = orig_shape
|
||||
x = x.reshape(B, T * H * W, D)
|
||||
|
||||
cx = self.cond_emb # (B, H*W, cond_emb_dim)
|
||||
# cond_emb は全モジュール共有の cx、depth_emb はこのモジュール用の depth vector。
|
||||
# ここで加算することで cx の N コピーを避ける。整数 layer_idx は参照せず
|
||||
# テンソル属性 depth_emb だけを足すので、torch.compile は単一グラフを維持する
|
||||
# (gradient checkpointing 下では cond_local は領域内で再計算され retain されない)。
|
||||
cond_local = self.cond_emb + self.depth_emb # (B, H*W, cond_emb_dim)
|
||||
|
||||
# CFG 推論用 (学習時は通らない想定)
|
||||
if x.shape[0] // 2 == cx.shape[0]:
|
||||
cx = cx.repeat(2, 1, 1)
|
||||
if x.shape[0] // 2 == cond_local.shape[0]:
|
||||
cond_local = cond_local.repeat(2, 1, 1)
|
||||
|
||||
# T=1 固定前提なので S == H*W のはず
|
||||
assert x.shape[1] == cx.shape[1], (
|
||||
f"LLLite seq mismatch ({self.lllite_name}): x={x.shape[1]} vs cond_emb={cx.shape[1]}"
|
||||
assert x.shape[1] == cond_local.shape[1], (
|
||||
f"LLLite seq mismatch ({self.lllite_name}): x={x.shape[1]} vs cond_emb={cond_local.shape[1]}"
|
||||
)
|
||||
|
||||
# depth embedding を加算 (zero-init なので学習序盤は cx のまま)
|
||||
if self._depth_embeds_ref:
|
||||
depth_e = self._depth_embeds_ref[0][self.layer_idx] # (cond_emb_dim,)
|
||||
cond_local = cx + depth_e # broadcast over (B, S)
|
||||
else:
|
||||
cond_local = cx
|
||||
|
||||
h = F.silu(self.down(x)) # (B, S, mlp)
|
||||
|
||||
# FiLM パラメータ (cond_local 由来、zero-init で identity)
|
||||
@@ -318,6 +324,8 @@ class ControlNetLLLiteDiT(nn.Module):
|
||||
cond_resblocks: int = 1,
|
||||
use_aspp: bool = False,
|
||||
aspp_dilations: Tuple[int, ...] = ASPP_DEFAULT_DILATIONS,
|
||||
cond_in_channels: int = 3,
|
||||
inpaint_masked_input: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
@@ -333,11 +341,16 @@ class ControlNetLLLiteDiT(nn.Module):
|
||||
self.cond_resblocks = cond_resblocks
|
||||
self.use_aspp = use_aspp
|
||||
self.aspp_dilations = tuple(aspp_dilations) if use_aspp else ()
|
||||
# 4ch (RGB+mask) inpainting 用の付加情報。inpaint_masked_input は学習側の RGB マスキング方針を
|
||||
# 記録するためのフラグで、モデル forward の挙動には影響しない (メタデータ復元用)。
|
||||
self.cond_in_channels = cond_in_channels
|
||||
self.inpaint_masked_input = inpaint_masked_input
|
||||
|
||||
# cond image (B,3,H*16,W*16) -> (B, S, cond_emb_dim)
|
||||
# cond image (B, cond_in_channels, H*16, W*16) -> (B, S, cond_emb_dim)
|
||||
self.conditioning1 = _Conditioning1(
|
||||
cond_dim, cond_emb_dim, cond_resblocks,
|
||||
use_aspp=use_aspp, aspp_dilations=aspp_dilations,
|
||||
cond_in_channels=cond_in_channels,
|
||||
)
|
||||
|
||||
modules = self._create_modules(dit, cond_emb_dim, mlp_dim, atomics, dropout, multiplier)
|
||||
@@ -348,14 +361,16 @@ class ControlNetLLLiteDiT(nn.Module):
|
||||
self.depth_embeds = nn.Parameter(torch.zeros(n, cond_emb_dim))
|
||||
for i, m in enumerate(self.lllite_modules):
|
||||
m.layer_idx = i
|
||||
m._depth_embeds_ref = [self.depth_embeds]
|
||||
|
||||
aspp_info = f"aspp={'on' + str(list(self.aspp_dilations)) if use_aspp else 'off'}"
|
||||
inpaint_info = (
|
||||
f", inpaint=on(masked_input={inpaint_masked_input})" if cond_in_channels == 4 else ""
|
||||
)
|
||||
logger.info(
|
||||
f"ControlNet-LLLite (Anima v{LLLITE_ARCH_VERSION}): created {n} modules for "
|
||||
f"target={target_layers!r} (atomics={list(atomics)}), "
|
||||
f"cond_dim={cond_dim}, cond_resblocks={cond_resblocks}, {aspp_info}, "
|
||||
f"cond_emb_dim={cond_emb_dim}, mlp_dim={mlp_dim}"
|
||||
f"cond_in_channels={cond_in_channels}, cond_dim={cond_dim}, cond_resblocks={cond_resblocks}, {aspp_info}, "
|
||||
f"cond_emb_dim={cond_emb_dim}, mlp_dim={mlp_dim}{inpaint_info}"
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -430,10 +445,21 @@ class ControlNetLLLiteDiT(nn.Module):
|
||||
if cond_image is None:
|
||||
for m in self.lllite_modules:
|
||||
m.cond_emb = None
|
||||
m.depth_emb = None
|
||||
return
|
||||
cx = self.conditioning1(cond_image) # (B, S, cond_emb_dim)
|
||||
for m in self.lllite_modules:
|
||||
m.cond_emb = cx
|
||||
# 共有の cx を全モジュールに同一テンソルとして持たせ (N コピーを避ける)、
|
||||
# depth embedding はこのモジュール用の (cond_emb_dim,) スライスだけを渡す。
|
||||
# 加算は forward 内で行う (cx の N 倍メモリを回避するため)。
|
||||
#
|
||||
# depth_embeds[layer_idx] の index 参照はここ (compile 領域の外) で毎ステップ
|
||||
# 行う。forward 内で整数 layer_idx を参照すると torch.compile がブロック毎に
|
||||
# 別グラフを焼くため、参照はここに残す。また index 参照を毎ステップ行うことで
|
||||
# SelectBackward が毎回張り直され、2 回目以降の backward でも depth_embeds へ
|
||||
# 正しく勾配が流れる (__init__ で一度だけ index すると graph 再利用で破綻する)。
|
||||
m.cond_emb = cx # 全モジュールで共有 (同一テンソル)
|
||||
m.depth_emb = self.depth_embeds[m.layer_idx] # (cond_emb_dim,), broadcast over (B, S)
|
||||
|
||||
def clear_cond_image(self):
|
||||
self.set_cond_image(None)
|
||||
@@ -478,6 +504,11 @@ class AnimaControlNetLLLiteWrapper(nn.Module):
|
||||
f"cond_image HW mismatch: latent={x.shape[-2]}x{x.shape[-1]} -> expected "
|
||||
f"{expected_h}x{expected_w}, got {cond_image.shape[-2]}x{cond_image.shape[-1]}"
|
||||
)
|
||||
expected_c = self.lllite.cond_in_channels
|
||||
assert cond_image.shape[1] == expected_c, (
|
||||
f"cond_image channel mismatch: expected {expected_c} (cond_in_channels), "
|
||||
f"got {cond_image.shape[1]}"
|
||||
)
|
||||
self.lllite.set_cond_image(cond_image)
|
||||
return self.dit(x, timesteps, context, **kwargs)
|
||||
|
||||
@@ -763,6 +794,41 @@ if __name__ == "__main__":
|
||||
assert torch.allclose(y_a, y_a_ref), "ASPP-on zero-init forward mismatch"
|
||||
logger.info(" ASPP-on zero-init forward OK")
|
||||
|
||||
# 4ch (inpainting) パス: 構築 + conv1 入力チャネル + zero-init forward + save/load round-trip
|
||||
dit_4ch = _DummyDiT(num_blocks=2, dim=64, ctx_dim=128)
|
||||
lllite_4ch = ControlNetLLLiteDiT(
|
||||
dit_4ch, cond_emb_dim=32, mlp_dim=64, target_layers="self_attn_q",
|
||||
cond_dim=64, cond_resblocks=1, cond_in_channels=4, inpaint_masked_input=True,
|
||||
)
|
||||
assert lllite_4ch.cond_in_channels == 4
|
||||
assert lllite_4ch.inpaint_masked_input is True
|
||||
assert lllite_4ch.conditioning1.conv1.in_channels == 4
|
||||
# 3ch 重みは 4ch モデルにそのまま load できないが、4ch round-trip は通る
|
||||
lllite_4ch.apply_to()
|
||||
wrapper_4ch = AnimaControlNetLLLiteWrapper(dit_4ch, lllite_4ch)
|
||||
H4, W4 = 8, 8
|
||||
cond4 = torch.randn(1, 4, H4 * 16, W4 * 16)
|
||||
wrapper_4ch.lllite.set_cond_image(cond4)
|
||||
cx4 = wrapper_4ch.lllite.lllite_modules[0].cond_emb
|
||||
assert cx4 is not None and cx4.shape == (1, H4 * W4, 32), f"4ch cond_emb shape: {cx4.shape}"
|
||||
# zero-init forward
|
||||
mod4 = wrapper_4ch.lllite.lllite_modules[0]
|
||||
x4 = torch.randn(1, H4 * W4, mod4.org_module[0].in_features)
|
||||
y4 = mod4(x4)
|
||||
y4_ref = mod4.org_forward(x4)
|
||||
assert torch.allclose(y4, y4_ref), "4ch zero-init forward mismatch"
|
||||
# Wrapper の cond_image チャネル assert: 3ch を渡すと AssertionError になる
|
||||
x_lat = torch.randn(1, 16, 1, H4 * 2, W4 * 2) # dummy latent
|
||||
try:
|
||||
# 3ch cond を渡すと拒否されるか
|
||||
wrapper_4ch(x_lat, torch.zeros(1), torch.zeros(1, 1, 1), cond_image=torch.randn(1, 3, H4 * 16, W4 * 16))
|
||||
raise AssertionError("expected channel mismatch assert")
|
||||
except AssertionError as e:
|
||||
msg = str(e)
|
||||
if "channel mismatch" not in msg and "expected 4" not in msg:
|
||||
raise
|
||||
logger.info(" 4ch (inpainting) build / forward / channel assert OK")
|
||||
|
||||
# 非デフォルト dilations
|
||||
lllite_dil = ControlNetLLLiteDiT(
|
||||
_DummyDiT(num_blocks=2, dim=64, ctx_dim=128),
|
||||
@@ -872,6 +938,36 @@ if __name__ == "__main__":
|
||||
assert torch.allclose(sd_orig[k].float(), sd_loaded[k].float()), f"mismatch at {k}"
|
||||
logger.info(" save / load round-trip OK")
|
||||
|
||||
# 4ch round-trip
|
||||
tmp4 = tmp + ".4ch.safetensors"
|
||||
try:
|
||||
meta4 = {
|
||||
"lllite.version": LLLITE_ARCH_VERSION,
|
||||
"lllite.cond_emb_dim": "32",
|
||||
"lllite.mlp_dim": "64",
|
||||
"lllite.target_layers": "self_attn_q",
|
||||
"lllite.cond_dim": "64",
|
||||
"lllite.cond_resblocks": "1",
|
||||
"lllite.cond_in_channels": "4",
|
||||
"lllite.inpaint_masked_input": "true",
|
||||
}
|
||||
save_lllite_model(tmp4, lllite_4ch, dtype=torch.float32, metadata=meta4)
|
||||
dit_4ch_b = _DummyDiT(num_blocks=2, dim=64, ctx_dim=128)
|
||||
lllite_4ch_b = ControlNetLLLiteDiT(
|
||||
dit_4ch_b, cond_emb_dim=32, mlp_dim=64, target_layers="self_attn_q",
|
||||
cond_dim=64, cond_resblocks=1, cond_in_channels=4, inpaint_masked_input=True,
|
||||
)
|
||||
load_lllite_weights(lllite_4ch_b, tmp4, strict=True)
|
||||
sd_a = lllite_4ch.state_dict()
|
||||
sd_b = lllite_4ch_b.state_dict()
|
||||
assert set(sd_a.keys()) == set(sd_b.keys())
|
||||
for k in sd_a:
|
||||
assert torch.allclose(sd_a[k].float(), sd_b[k].float()), f"4ch round-trip mismatch at {k}"
|
||||
logger.info(" 4ch save / load round-trip OK")
|
||||
finally:
|
||||
if os.path.exists(tmp4):
|
||||
os.unlink(tmp4)
|
||||
|
||||
# 旧形式 (lllite_modules.* キー) は reject される
|
||||
legacy_sd = {"lllite_modules.0.up.weight": torch.zeros(1)}
|
||||
from safetensors.torch import save_file as _save_legacy
|
||||
|
||||
Reference in New Issue
Block a user