| 57 | |
| 58 | |
| 59 | class VQGANConfig(PretrainedConfig): |
| 60 | model_type = "vqgan" |
| 61 | |
| 62 | def __init__( |
| 63 | self, |
| 64 | resolution=256, |
| 65 | num_channels=3, |
| 66 | hidden_channels=128, |
| 67 | channel_mult=(1, 2, 2, 4, 6), |
| 68 | num_res_blocks=2, |
| 69 | attn_resolutions=(), |
| 70 | no_attn_mid_block=True, |
| 71 | z_channels=64, |
| 72 | num_embeddings=8192, |
| 73 | quantized_embed_dim=64, |
| 74 | dropout=0.0, |
| 75 | resample_with_conv=True, |
| 76 | commitment_cost=0.25 |
| 77 | ): |
| 78 | self.resolution = resolution |
| 79 | self.num_channels = num_channels |
| 80 | self.hidden_channels = hidden_channels |
| 81 | self.channel_mult = channel_mult |
| 82 | self.num_res_blocks = num_res_blocks |
| 83 | self.attn_resolutions = attn_resolutions |
| 84 | self.no_attn_mid_block = no_attn_mid_block |
| 85 | self.z_channels = z_channels |
| 86 | self.num_embeddings = num_embeddings |
| 87 | self.quantized_embed_dim = quantized_embed_dim |
| 88 | self.dropout = dropout |
| 89 | self.resample_with_conv = resample_with_conv |
| 90 | self.commitment_cost = commitment_cost |
| 91 | |
| 92 | @classmethod |
| 93 | def get_default_config(cls, updates=None): |
| 94 | config = function_args_to_config(cls.__init__) |
| 95 | if updates is not None: |
| 96 | config.update(ConfigDict(updates).copy_and_resolve_references()) |
| 97 | config.num_resolutions = len(config.channel_mult) |
| 98 | return config |
| 99 | |
| 100 | @classmethod |
| 101 | def load_config(cls, path): |
| 102 | return cls.get_default_config(cls) |
| 103 | |
| 104 | |
| 105 | class VQGANModel(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected