File size: 2,186 Bytes
1a18f22 | 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 | """build_backbone: instantiate one of {jit, pixelgen, deco, pixeldit} pixel-space
denoisers for mask-concat conditioning. Each returns a net callable as net(x, t, y)
-> (N, C>=img_ch, H, W); the caller slices [:, :img_ch] (backbone-agnostic decouple).
in_channels = img_channels + cond_channels; a single dummy class (num_classes=1).
Unified 'Base' tier (~130-150M) for the P1 backbone bake-off (P1 = native arch under a
common flow-matching objective; perceptual/DCT/FD losses are P2 levers, not used here)."""
import os
import sys
_SOTA = "/home/wzhang/LSC/Code/NPJ/sota"
def _add(path):
if path not in sys.path:
sys.path.insert(0, path)
def build_backbone(backbone: str, model_name: str, img_size: int,
in_channels: int, num_classes: int = 1):
bk = backbone.lower()
if bk == "jit":
_add(os.path.join(_SOTA, "JiT"))
from model_jit import JiT_models
return JiT_models[model_name](input_size=img_size, in_channels=in_channels,
num_classes=num_classes)
if bk == "pixelgen":
_add(os.path.join(_SOTA, "PixelGen", "src", "models", "transformer"))
import importlib
jit = importlib.import_module("JiT") # PixelGen's self-contained JiT.py
return jit.JiT_models[model_name](input_size=img_size, in_channels=in_channels,
num_classes=num_classes)
if bk == "deco":
_add(os.path.join(_SOTA, "DeCo", "src", "models", "transformer"))
from dit_c2i_DeCo import PixNerDiT
return PixNerDiT(in_channels=in_channels, patch_size=16, num_groups=12,
hidden_size=768, hidden_size_x=32, num_blocks=13,
num_cond_blocks=12, num_classes=num_classes)
if bk == "pixeldit":
_add(os.path.join(_SOTA, "PixelDiT"))
from pixdit_core.pixeldit_c2i import PixDiT
return PixDiT(in_channels=in_channels, num_groups=10, hidden_size=640,
pixel_hidden_size=16, patch_depth=9, pixel_depth=4,
patch_size=16, num_classes=num_classes)
raise ValueError(f"unknown backbone: {backbone}")
|