first commit

lib/attention.py

@@ -16,7 +16,7 @@ from torch import nn, einsum
 from einops import rearrange, repeat
 from typing import Optional, Any
 
-from utils import checkpoint
+from ..utils import checkpoint
 
 try:
     import xformers
@@ -351,4 +351,4 @@ class SpatialTransformer(nn.Module):
         x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
         if not self.use_linear:
             x = self.proj_out(x)
-        return x + x_in
+        return x + x_in

lib/ddpm_multi.py

@@ -30,7 +30,7 @@ from torchvision.utils import make_grid
 from pytorch_lightning.utilities.distributed import rank_zero_only
 from omegaconf import ListConfig
 
-from utils import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
+from ..utils import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config
 from lib.distributions import normal_kl, DiagonalGaussianDistribution
 from lib.autoencoder import IdentityFirstStage, AutoencoderKL
 from lib.util import make_beta_schedule, extract_into_tensor, noise_like
@@ -1798,4 +1798,4 @@ class LatentUpscaleFinetuneDiffusion(LatentFinetuneDiffusion):
     def log_images(self, *args, **kwargs):
         log = super().log_images(*args, **kwargs)
         log["lr"] = rearrange(args[0]["lr"], 'b h w c -> b c h w')
-        return log
+        return log

lib/openaimodel.py

@@ -26,7 +26,7 @@ from lib.util import (
     timestep_embedding,
 )
 from attention import SpatialTransformer
-from utils import exists
+from ..utils import exists
 
 
 # dummy replace
@@ -793,4 +793,4 @@ class UNetModel(nn.Module):
         if self.predict_codebook_ids:
             return self.id_predictor(h)
         else:
-            return self.out(h)
+            return self.out(h)

lib/util.py

@@ -25,7 +25,7 @@ import torch.nn as nn
 import numpy as np
 from einops import repeat
 
-from utils import instantiate_from_config
+from ..utils import instantiate_from_config
 
 
 def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
@@ -277,4 +277,4 @@ class HybridConditioner(nn.Module):
 def noise_like(shape, device, repeat=False):
     repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
     noise = lambda: torch.randn(shape, device=device)
-    return repeat_noise() if repeat else noise()
+    return repeat_noise() if repeat else noise()