(
self,
# ------------------- DEFAULTS FROM Stable Diffusion 2-1
image_size=32,
in_channels=4,
out_channels=4,
model_channels=320,
attention_resolutions=[ 4, 2, 1 ],
num_res_blocks=2,
channel_mult=[ 1, 2, 4, 4 ],
num_heads=-1,
num_head_channels=64,
use_spatial_transformer=True, # custom transformer support
use_linear_in_transformer=True,
transformer_depth=1, # custom transformer support
context_dim=1024, # custom transformer support
legacy=False,
load_from_ckpt=None,
# ----------------------------------------------------
dropout=0,
conv_resample=True,
dims=2,
num_classes=None,
use_checkpoint=False,
use_fp16=False,
use_bf16=False,
num_heads_upsample=-1,
use_scale_shift_norm=False,
resblock_updown=False,
use_new_attention_order=False,
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
disable_self_attentions=None,
num_attention_blocks=None,
disable_middle_self_attn=False,
adm_in_channels=None,
concat_context=True,
)
| 452 | increased efficiency. |
| 453 | """ |
| 454 | def __init__( |
| 455 | self, |
| 456 | # ------------------- DEFAULTS FROM Stable Diffusion 2-1 |
| 457 | image_size=32, |
| 458 | in_channels=4, |
| 459 | out_channels=4, |
| 460 | model_channels=320, |
| 461 | attention_resolutions=[ 4, 2, 1 ], |
| 462 | num_res_blocks=2, |
| 463 | channel_mult=[ 1, 2, 4, 4 ], |
| 464 | num_heads=-1, |
| 465 | num_head_channels=64, |
| 466 | use_spatial_transformer=True, # custom transformer support |
| 467 | use_linear_in_transformer=True, |
| 468 | transformer_depth=1, # custom transformer support |
| 469 | context_dim=1024, # custom transformer support |
| 470 | legacy=False, |
| 471 | load_from_ckpt=None, |
| 472 | # ---------------------------------------------------- |
| 473 | dropout=0, |
| 474 | conv_resample=True, |
| 475 | dims=2, |
| 476 | num_classes=None, |
| 477 | use_checkpoint=False, |
| 478 | use_fp16=False, |
| 479 | use_bf16=False, |
| 480 | num_heads_upsample=-1, |
| 481 | use_scale_shift_norm=False, |
| 482 | resblock_updown=False, |
| 483 | use_new_attention_order=False, |
| 484 | n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model |
| 485 | disable_self_attentions=None, |
| 486 | num_attention_blocks=None, |
| 487 | disable_middle_self_attn=False, |
| 488 | adm_in_channels=None, |
| 489 | concat_context=True, |
| 490 | ): |
| 491 | super().__init__() |
| 492 | print(f'[OpenAIUNet] Creating Model') |
| 493 | if use_spatial_transformer: |
| 494 | assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' |
| 495 | |
| 496 | if context_dim is not None: |
| 497 | assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' |
| 498 | from omegaconf.listconfig import ListConfig |
| 499 | if type(context_dim) == ListConfig: |
| 500 | context_dim = list(context_dim) |
| 501 | |
| 502 | if num_heads_upsample == -1: |
| 503 | num_heads_upsample = num_heads |
| 504 | |
| 505 | if num_heads == -1: |
| 506 | assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' |
| 507 | |
| 508 | if num_head_channels == -1: |
| 509 | assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' |
| 510 | |
| 511 | self.image_size = image_size |
nothing calls this directly
no test coverage detected