(
self,
*,
in_chans: int = 256,
out_chans: int = 3,
upscale_factor: int = 16,
# ViT params
pos_embed_rope_base: float = 100.0,
pos_embed_rope_min_period: Optional[float] = None,
pos_embed_rope_max_period: Optional[float] = None,
pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate",
pos_embed_rope_shift_coords: Optional[float] = None,
pos_embed_rope_jitter_coords: Optional[float] = None,
pos_embed_rope_rescale_coords: Optional[float] = None,
pos_embed_rope_dtype: str = "bf16",
embed_dim: int = 1024,
depth: int = 24,
num_heads: int = 16,
ffn_ratio: float = 4.0,
qkv_bias: bool = True,
drop_path_rate: float = 0.0,
layerscale_init: Optional[float] = None,
norm_layer: str = "layernorm",
ffn_layer: str = "swiglu",
ffn_bias: bool = True,
proj_bias: bool = True,
mask_k_bias: bool = False,
device: Optional[Any] = None,
use_qk_norm: bool = False,
**ignored_kwargs,
)
| 20 | """ |
| 21 | |
| 22 | def __init__( |
| 23 | self, |
| 24 | *, |
| 25 | in_chans: int = 256, |
| 26 | out_chans: int = 3, |
| 27 | upscale_factor: int = 16, |
| 28 | # ViT params |
| 29 | pos_embed_rope_base: float = 100.0, |
| 30 | pos_embed_rope_min_period: Optional[float] = None, |
| 31 | pos_embed_rope_max_period: Optional[float] = None, |
| 32 | pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate", |
| 33 | pos_embed_rope_shift_coords: Optional[float] = None, |
| 34 | pos_embed_rope_jitter_coords: Optional[float] = None, |
| 35 | pos_embed_rope_rescale_coords: Optional[float] = None, |
| 36 | pos_embed_rope_dtype: str = "bf16", |
| 37 | embed_dim: int = 1024, |
| 38 | depth: int = 24, |
| 39 | num_heads: int = 16, |
| 40 | ffn_ratio: float = 4.0, |
| 41 | qkv_bias: bool = True, |
| 42 | drop_path_rate: float = 0.0, |
| 43 | layerscale_init: Optional[float] = None, |
| 44 | norm_layer: str = "layernorm", |
| 45 | ffn_layer: str = "swiglu", |
| 46 | ffn_bias: bool = True, |
| 47 | proj_bias: bool = True, |
| 48 | mask_k_bias: bool = False, |
| 49 | device: Optional[Any] = None, |
| 50 | use_qk_norm: bool = False, |
| 51 | **ignored_kwargs, |
| 52 | ): |
| 53 | super().__init__() |
| 54 | if len(ignored_kwargs) > 0: |
| 55 | logger.warning(f"Ignored kwargs: {ignored_kwargs}") |
| 56 | del ignored_kwargs |
| 57 | |
| 58 | norm_layer_cls = norm_layer_dict[norm_layer] |
| 59 | self.embed_dim = embed_dim |
| 60 | self.num_heads = num_heads |
| 61 | |
| 62 | # 1. Input projection |
| 63 | self.proj_in = nn.Conv2d( |
| 64 | in_chans, embed_dim, kernel_size=1, bias=proj_bias |
| 65 | ) |
| 66 | |
| 67 | # 2. RoPE |
| 68 | self.rope_embed = RopePositionEmbedding( |
| 69 | embed_dim=embed_dim, |
| 70 | num_heads=num_heads, |
| 71 | base=pos_embed_rope_base, |
| 72 | min_period=pos_embed_rope_min_period, |
| 73 | max_period=pos_embed_rope_max_period, |
| 74 | normalize_coords=pos_embed_rope_normalize_coords, |
| 75 | shift_coords=pos_embed_rope_shift_coords, |
| 76 | jitter_coords=pos_embed_rope_jitter_coords, |
| 77 | rescale_coords=pos_embed_rope_rescale_coords, |
| 78 | dtype=dtype_dict[pos_embed_rope_dtype], |
| 79 | device=device, |
nothing calls this directly
no test coverage detected