| 41 | |
| 42 | |
| 43 | class MVDiffusion(pl.LightningModule): |
| 44 | def __init__( |
| 45 | self, |
| 46 | stable_diffusion_config, |
| 47 | drop_cond_prob=0.1, |
| 48 | ): |
| 49 | super(MVDiffusion, self).__init__() |
| 50 | |
| 51 | self.drop_cond_prob = drop_cond_prob |
| 52 | |
| 53 | self.register_schedule() |
| 54 | |
| 55 | # init modules |
| 56 | pipeline = DiffusionPipeline.from_pretrained(**stable_diffusion_config) |
| 57 | pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( |
| 58 | pipeline.scheduler.config, timestep_spacing='trailing' |
| 59 | ) |
| 60 | self.pipeline = pipeline |
| 61 | |
| 62 | train_sched = DDPMScheduler.from_config(self.pipeline.scheduler.config) |
| 63 | if isinstance(self.pipeline.unet, UNet2DConditionModel): |
| 64 | self.pipeline.unet = RefOnlyNoisedUNet(self.pipeline.unet, train_sched, self.pipeline.scheduler) |
| 65 | |
| 66 | self.train_scheduler = train_sched # use ddpm scheduler during training |
| 67 | |
| 68 | self.unet = pipeline.unet |
| 69 | |
| 70 | # validation output buffer |
| 71 | self.validation_step_outputs = [] |
| 72 | |
| 73 | def register_schedule(self): |
| 74 | self.num_timesteps = 1000 |
| 75 | |
| 76 | # replace scaled_linear schedule with linear schedule as Zero123++ |
| 77 | beta_start = 0.00085 |
| 78 | beta_end = 0.0120 |
| 79 | betas = torch.linspace(beta_start, beta_end, 1000, dtype=torch.float32) |
| 80 | |
| 81 | alphas = 1. - betas |
| 82 | alphas_cumprod = torch.cumprod(alphas, dim=0) |
| 83 | alphas_cumprod_prev = torch.cat([torch.ones(1, dtype=torch.float64), alphas_cumprod[:-1]], 0) |
| 84 | |
| 85 | self.register_buffer('betas', betas.float()) |
| 86 | self.register_buffer('alphas_cumprod', alphas_cumprod.float()) |
| 87 | self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float()) |
| 88 | |
| 89 | # calculations for diffusion q(x_t | x_{t-1}) and others |
| 90 | self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod).float()) |
| 91 | self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1 - alphas_cumprod).float()) |
| 92 | |
| 93 | self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod).float()) |
| 94 | self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1).float()) |
| 95 | |
| 96 | def on_fit_start(self): |
| 97 | device = torch.device(f'cuda:{self.global_rank}') |
| 98 | self.pipeline.to(device) |
| 99 | if self.global_rank == 0: |
| 100 | os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True) |
nothing calls this directly
no outgoing calls
no test coverage detected