| 12 | |
| 13 | |
| 14 | class JoDiffusionModel(ModelMixin, ConfigMixin): |
| 15 | |
| 16 | @register_to_config |
| 17 | def __init__( |
| 18 | self, |
| 19 | text_dim: int = 768, |
| 20 | inner_text_dim: int = 64, |
| 21 | clip_img_dim: int = 512, |
| 22 | num_text_tokens: int = 77, |
| 23 | num_attention_heads: int = 24, |
| 24 | attention_head_dim: int = 64, |
| 25 | in_channels: int = 4, |
| 26 | out_channels: int = 4, |
| 27 | num_layers: int = 30, |
| 28 | dropout: float = 0.0, |
| 29 | norm_num_groups: int = 32, |
| 30 | cross_attention_dim: Optional[int] = None, |
| 31 | attention_bias: bool = False, |
| 32 | sample_size: int = 64, |
| 33 | num_vector_embeds: Optional[int] = None, |
| 34 | patch_size: int = 2, |
| 35 | activation_fn: str = "gelu", |
| 36 | num_embeds_ada_norm: int = 1000, |
| 37 | use_linear_projection: bool = False, |
| 38 | only_cross_attention: bool = False, |
| 39 | upcast_attention: bool = False, |
| 40 | norm_type: str = "layer_norm", |
| 41 | block_type: str = "unidiffuser", |
| 42 | pre_layer_norm: bool = False, |
| 43 | use_timestep_embedding: bool = False, |
| 44 | norm_elementwise_affine: bool = True, |
| 45 | use_patch_pos_embed: bool = False, |
| 46 | ff_final_dropout: bool = True, |
| 47 | use_data_type_embedding: bool = True, |
| 48 | ): |
| 49 | super().__init__() |
| 50 | |
| 51 | # 0. Handle dimensions |
| 52 | self.inner_dim = num_attention_heads * attention_head_dim |
| 53 | |
| 54 | assert sample_size is not None, "UniDiffuserModel over patched input must provide sample_size" |
| 55 | self.sample_size = sample_size |
| 56 | self.in_channels = in_channels |
| 57 | self.out_channels = in_channels if out_channels is None else out_channels |
| 58 | |
| 59 | self.patch_size = patch_size |
| 60 | # Assume image is square... |
| 61 | self.num_patches = (self.sample_size // patch_size) * (self.sample_size // patch_size) |
| 62 | |
| 63 | # 1. Define input layers |
| 64 | # 1.1 Input layers for label and image input |
| 65 | # For now, only support patch input for VAE latent image input |
| 66 | self.pre_text = nn.Linear(text_dim, inner_text_dim) |
| 67 | self.text_in = nn.Linear(inner_text_dim, self.inner_dim) |
| 68 | self.vae_img_in = PatchEmbed( |
| 69 | height=sample_size, |
| 70 | width=sample_size, |
| 71 | patch_size=patch_size, |