cleanup, remove non-essential stuff

Author: brianfitzgerald

cn_eval.out.538931

@@ -58,7 +58,6 @@ def __init__(
             embed_dim=self.hidden_size,
             strict_img_size=False,
         )
-        self.is_8b = True
 
     def forward(
         self,
@@ -67,9 +66,10 @@ def forward(
         y: Tensor,
         scale: int = 1,
         timestep: Optional[Tensor] = None,
+        is_8b: bool = False
     ) -> Tuple[Tensor, List[Tensor]]:
 
-        if not self.is_8b:
+        if not is_8b:
             x = self.x_embedder(x)
         timestep = timestep * 1000
         c = self.t_embedder(timestep, dtype=x.dtype)
@@ -83,7 +83,7 @@ def forward(
 
         for block in self.transformer_blocks:
             out = block(x, c)
-            if self.is_8b:
+            if is_8b:
                 x = out
             block_out += (out,)
 

dit_embedder.py

@@ -159,21 +159,17 @@ def apply_model(self, x, sigma, c_crossattn=None, y=None, skip_layers=[], contro
             controlnet_cond = controlnet_cond.to(dtype=x.dtype, device=x.device)
             controlnet_cond = controlnet_cond.repeat(x.shape[0], 1, 1, 1)
 
-            # Some ControlNets don't use the y_cond input, so we need to check if it's needed.
-            if y_cond.shape[-1] != self.control_model.y_embedder.mlp[0].in_features:
+            # 8B ControlNets were trained with a slightly different architecture.
+            is_8b = y_cond.shape[-1] == self.control_model.y_embedder.mlp[0].in_features
+            if not is_8b:
                 y_cond = self.diffusion_model.y_embedder(y)
-            hw = x.shape[-2:]
 
             x_controlnet = x
-            # HACK
-            # x_controlnet = torch.load("/weka/home-brianf/x_8b.pt")
-            # controlnet_cond = torch.load("/weka/home-brianf/x_cond_8b.pt")
-            # y_cond = torch.load("/weka/home-brianf/y_cond_8b.pt")
-            if self.control_model.is_8b:
+            if is_8b:
+                hw = x.shape[-2:]
                 x_controlnet = self.diffusion_model.x_embedder(x) + self.diffusion_model.cropped_pos_embed(hw)
-                # y_cond[0] = torch.zeros_like(y_cond[0])
             controlnet_hidden_states = self.control_model(
-                x_controlnet, controlnet_cond, y_cond, 1, sigma.to(torch.float32)
+                x_controlnet, controlnet_cond, y_cond, 1, sigma.to(torch.float32), is_8b
             )
         model_output = self.diffusion_model(
             x.to(dtype),
@@ -747,72 +743,3 @@ def encode(self, image):
         std = torch.exp(0.5 * logvar)
         return mean + std * torch.randn_like(mean)
 
-
-class DiagonalGaussianDistribution:
-    def __init__(self, parameters, deterministic=False, chunk_dim: int = 1):
-        self.parameters = parameters
-        self.mean, self.logvar = torch.chunk(parameters, 2, dim=chunk_dim)
-        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
-        self.deterministic = deterministic
-        self.std = torch.exp(0.5 * self.logvar)
-        self.var = torch.exp(self.logvar)
-        if self.deterministic:
-            self.var = self.std = torch.zeros_like(self.mean).to(
-                device=self.parameters.device
-            )
-
-    def sample(self):
-        x = self.mean + self.std * torch.randn(self.mean.shape).to(
-            device=self.parameters.device
-        )
-        return x
-
-    def kl(self, other=None):
-        if self.deterministic:
-            return torch.Tensor([0.0])
-        else:
-            if other is None:
-                return 0.5 * torch.sum(
-                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
-                    dim=list(range(1, self.mean.ndim)),
-                )
-            else:
-                return 0.5 * torch.sum(
-                    torch.pow(self.mean - other.mean, 2) / other.var
-                    + self.var / other.var
-                    - 1.0
-                    - self.logvar
-                    + other.logvar,
-                    dim=list(range(1, self.mean.ndim)),
-                )
-
-    def nll(self, sample, dims=[1, 2, 3]):
-        if self.deterministic:
-            return torch.Tensor([0.0])
-        logtwopi = np.log(2.0 * np.pi)
-        return 0.5 * torch.sum(
-            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
-            dim=dims,
-        )
-
-    def mode(self):
-        return self.mean
-
-
-class DiagonalGaussianRegularizer(nn.Module):
-    def __init__(self, sample: bool = True, chunk_dim: int = 1):
-        super().__init__()
-        self.sample = sample
-        self.chunk_dim = chunk_dim
-
-    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
-        log = dict()
-        posterior = DiagonalGaussianDistribution(z, chunk_dim=self.chunk_dim)
-        if self.sample:
-            z = posterior.sample()
-        else:
-            z = posterior.mode()
-        kl_loss = posterior.kl()
-        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
-        log["kl_loss"] = kl_loss
-        return z, log

evaluate.py

@@ -24,7 +24,6 @@
     SDVAE,
     BaseModel,
     CFGDenoiser,
-    DiagonalGaussianRegularizer,
     SD3LatentFormat,
     SkipLayerCFGDenoiser,
 )
@@ -393,17 +392,8 @@ def vae_encode(self, image, controlnet_cond: bool = False) -> torch.Tensor:
         self.print("Encoded")
         return latent
 
-    def vae_encode_pkl(self, pkl_location: str) -> torch.Tensor:
-        with open(pkl_location, "rb") as f:
-            data = pickle.load(f)
-        latent = data["vae_f8_ch16.cond.sft.latent"]
-        latent, _ = DiagonalGaussianRegularizer()(latent)
-        latent = SD3LatentFormat().process_in(latent)
-        return latent
-
     def vae_encode_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
         tensor = tensor.unsqueeze(0)
-        latent, _ = DiagonalGaussianRegularizer()(tensor)
         latent = SD3LatentFormat().process_in(latent)
         return latent
 
@@ -454,7 +444,6 @@ def gen_image(
             controlnet_cond = self._image_to_latent(
                 controlnet_cond_image, width, height, True
             )
-            # controlnet_cond = self.vae_encode_pkl("/weka/home-brianf/controlnet_val/canny_8_3/pkl/data_6.pkl")
         neg_cond = self.get_cond("")
         seed_num = None
         pbar = tqdm(enumerate(prompts), position=0, leave=True)