main class
| 422 | |
| 423 | |
| 424 | class LatentDiffusion(DDPM): |
| 425 | """main class""" |
| 426 | def __init__(self, |
| 427 | first_stage_config, |
| 428 | cond_stage_config, |
| 429 | num_timesteps_cond=None, |
| 430 | cond_stage_key="caption", |
| 431 | cond_stage_trainable=False, |
| 432 | cond_stage_forward=None, |
| 433 | conditioning_key=None, |
| 434 | uncond_prob=0.2, |
| 435 | uncond_type="empty_seq", |
| 436 | scale_factor=1.0, |
| 437 | scale_by_std=False, |
| 438 | # added for LVDM |
| 439 | encoder_type="2d", |
| 440 | frame_cond=None, |
| 441 | only_model=False, |
| 442 | logdir=None, |
| 443 | empty_params_only=False, |
| 444 | *args, **kwargs): |
| 445 | self.num_timesteps_cond = default(num_timesteps_cond, 1) |
| 446 | self.scale_by_std = scale_by_std |
| 447 | assert self.num_timesteps_cond <= kwargs['timesteps'] |
| 448 | # for backwards compatibility after implementation of DiffusionWrapper |
| 449 | ckpt_path = kwargs.pop("ckpt_path", None) |
| 450 | ignore_keys = kwargs.pop("ignore_keys", []) |
| 451 | conditioning_key = default(conditioning_key, 'crossattn') |
| 452 | super().__init__(conditioning_key=conditioning_key, *args, **kwargs) |
| 453 | |
| 454 | self.cond_stage_trainable = cond_stage_trainable |
| 455 | self.cond_stage_key = cond_stage_key |
| 456 | self.empty_params_only = empty_params_only |
| 457 | try: |
| 458 | self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 |
| 459 | except: |
| 460 | self.num_downs = 0 |
| 461 | if not scale_by_std: |
| 462 | self.scale_factor = scale_factor |
| 463 | else: |
| 464 | self.register_buffer('scale_factor', torch.tensor(scale_factor)) |
| 465 | self.instantiate_first_stage(first_stage_config) |
| 466 | self.instantiate_cond_stage(cond_stage_config) |
| 467 | self.first_stage_config = first_stage_config |
| 468 | self.cond_stage_config = cond_stage_config |
| 469 | self.clip_denoised = False |
| 470 | |
| 471 | self.cond_stage_forward = cond_stage_forward |
| 472 | self.encoder_type = encoder_type |
| 473 | assert(encoder_type in ["2d", "3d"]) |
| 474 | self.uncond_prob = uncond_prob |
| 475 | self.classifier_free_guidance = True if uncond_prob > 0 else False |
| 476 | assert(uncond_type in ["zero_embed", "empty_seq"]) |
| 477 | self.uncond_type = uncond_type |
| 478 | |
| 479 | ## future frame prediction |
| 480 | self.frame_cond = frame_cond |
| 481 | if self.frame_cond: |
nothing calls this directly
no outgoing calls
no test coverage detected