(self,
first_stage_config,
cond_stage_config,
num_timesteps_cond=None,
cond_stage_key="image",
cond_stage_trainable=False,
concat_mode=True,
cond_stage_forward=None,
conditioning_key=None,
scale_factor=1.0,
scale_by_std=False,
force_null_conditioning=False,
*args, **kwargs)
| 531 | """main class""" |
| 532 | |
| 533 | def __init__(self, |
| 534 | first_stage_config, |
| 535 | cond_stage_config, |
| 536 | num_timesteps_cond=None, |
| 537 | cond_stage_key="image", |
| 538 | cond_stage_trainable=False, |
| 539 | concat_mode=True, |
| 540 | cond_stage_forward=None, |
| 541 | conditioning_key=None, |
| 542 | scale_factor=1.0, |
| 543 | scale_by_std=False, |
| 544 | force_null_conditioning=False, |
| 545 | *args, **kwargs): |
| 546 | self.force_null_conditioning = force_null_conditioning |
| 547 | self.num_timesteps_cond = default(num_timesteps_cond, 1) |
| 548 | self.scale_by_std = scale_by_std |
| 549 | assert self.num_timesteps_cond <= kwargs['timesteps'] |
| 550 | # for backwards compatibility after implementation of DiffusionWrapper |
| 551 | if conditioning_key is None: |
| 552 | conditioning_key = 'concat' if concat_mode else 'crossattn' |
| 553 | if cond_stage_config == '__is_unconditional__' and not self.force_null_conditioning: |
| 554 | conditioning_key = None |
| 555 | ckpt_path = kwargs.pop("ckpt_path", None) |
| 556 | reset_ema = kwargs.pop("reset_ema", False) |
| 557 | reset_num_ema_updates = kwargs.pop("reset_num_ema_updates", False) |
| 558 | ignore_keys = kwargs.pop("ignore_keys", []) |
| 559 | super().__init__(conditioning_key=conditioning_key, *args, **kwargs) |
| 560 | self.concat_mode = concat_mode |
| 561 | self.cond_stage_trainable = cond_stage_trainable |
| 562 | self.cond_stage_key = cond_stage_key |
| 563 | try: |
| 564 | self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 |
| 565 | except: |
| 566 | self.num_downs = 0 |
| 567 | if not scale_by_std: |
| 568 | self.scale_factor = scale_factor |
| 569 | else: |
| 570 | self.register_buffer('scale_factor', torch.tensor(scale_factor)) |
| 571 | self.instantiate_first_stage(first_stage_config) |
| 572 | self.instantiate_cond_stage(cond_stage_config) |
| 573 | self.cond_stage_forward = cond_stage_forward |
| 574 | self.clip_denoised = False |
| 575 | self.bbox_tokenizer = None |
| 576 | |
| 577 | self.restarted_from_ckpt = False |
| 578 | if ckpt_path is not None: |
| 579 | self.init_from_ckpt(ckpt_path, ignore_keys) |
| 580 | self.restarted_from_ckpt = True |
| 581 | if reset_ema: |
| 582 | assert self.use_ema |
| 583 | print( |
| 584 | f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint.") |
| 585 | self.model_ema = LitEma(self.model) |
| 586 | if reset_num_ema_updates: |
| 587 | print(" +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ ") |
| 588 | assert self.use_ema |
| 589 | self.model_ema.reset_num_updates() |
| 590 |
nothing calls this directly
no test coverage detected