Class Bokehlicious

method/model.py:12–318

from method.nn_util import PatchEmbedIR, PatchUnEmbedIR, IdentityMod, SkipConnection

class Bokehlicious(Module):

    def __init__(self,
                 # General Args
                 in_chans=3, img_range=1., dynamic_conv_k=4,
                 # In Stage Args
                 in_stage_use_pos_map: bool = True,
                 in_stage_use_bokeh_strength_map: bool = False,
                 # UNet Args
                 u_width=32,
                 u_depth=2,
                 u_block_config=None,
                 u_skip_connections=None,
                 # Enc Args
                 enc_blk_nums: List[int] = None,
                 enc_blks_use_pos_map: List[bool] = None,
                 # Deep Feature Extractor args
                 embed_dims=None,
                 depths=None,
                 num_heads=None,
                 init_values=None,
                 heads_ranges=None,
                 mlp_ratios=None,
                 drop_path_rate=0.1,
                 norm_layer=LayerNorm,
                 patch_norm=True,
                 use_checkpoints=None,
                 chunkwise_recurrents=None,
                 layerscales=None,
                 layer_init_values=1e-6,
                 positional_dfe=False,
                 use_dfe_norm_layer=True,
                 positional_conv_last=False,
                 # Dec Args
                 dec_blk_nums: List[int] = None,
                 dec_blks_use_pos_map: List[bool] = None,
                 # Out Stage Args
                 out_stage_use_pos_map: bool = True
                 ):
        super().__init__()

        self.in_chans = in_chans
        self.img_range = img_range

        # In Stage Args
        self.in_stage_use_pos_map = in_stage_use_pos_map
        self.in_stage_use_bokeh_strength_map = in_stage_use_bokeh_strength_map

        self.u_width = u_width
        self.u_depth = u_depth
        self.u_block_config = u_block_config or {'dw_expand': 1., 'ffn_expand': 2.,
                                                 'drop_out_rate': 0, 'attention_type': 'CA',
                                                 'activation_type': 'GELU', 'kernel_size': 3,
                                                 'inverted_conv': True,
                                                 }

        self.u_skip_connections = u_skip_connections or [True for _ in range(u_depth)]