incorporated new EMA into DIP

Author: dmarx

src/pytti/image_models/deep_image_prior.py

@@ -8,7 +8,9 @@
 import torch
 from torch import nn
 from torchvision.transforms import functional as TF
-from pytti.image_models import DifferentiableImage, EMAImage
+
+# from pytti.image_models import DifferentiableImage
+from pytti.image_models.ema import EMAImage, EMAParametersDict
 from PIL import Image
 from torch.nn import functional as F
 
@@ -44,8 +46,8 @@ def load_dip(input_depth, num_scales, offset_type, offset_groups, device):
     return dip_net
 
 
-# class DeepImagePrior(EMAImage):
-class DeepImagePrior(DifferentiableImage):
+class DeepImagePrior(EMAImage):
+    # class DeepImagePrior(DifferentiableImage):
     """
     https://github.com/nousr/deep-image-prior/
     """
@@ -69,7 +71,14 @@ def __init__(
         device="cuda",
         **kwargs,
     ):
-        super().__init__(width * scale, height * scale)
+        # super(super(EMAImage)).__init__()
+        nn.Module.__init__(self)
+        super().__init__(
+            width=width * scale,
+            height=height * scale,
+            decay=ema_val,
+            device=device,
+        )
         net = load_dip(
             input_depth=input_depth,
             num_scales=num_scales,
@@ -85,20 +94,38 @@ def __init__(
         # z = torch.cat(get_non_offset_params(net), get_offset_params(net))
         # logger.debug(z.shape)
         # super().__init__(width * scale, height * scale, z, ema_val)
-        self.net = net
+        # self.net = net
         # self.tensor = self.net.params()
         self.output_axes = ("n", "s", "y", "x")
         self.scale = scale
         self.device = device
 
-        self._net_input = torch.randn([1, input_depth, width, height], device=device)
+        # self._net_input = torch.randn([1, input_depth, width, height], device=device)
 
         self.lr = lr
         self.offset_lr_fac = offset_lr_fac
         # self._params = [
         #    {'params': get_non_offset_params(net), 'lr': lr},
         #    {'params': get_offset_params(net), 'lr': lr * offset_lr_fac}
         # ]
+        # z = {
+        #    'non_offset':get_non_offset_params(net),
+        #    'offset':get_offset_params(net),
+        # }
+        self.net = net
+        self._net_input = torch.randn([1, input_depth, width, height], device=device)
+
+        self.image_representation_parameters = EMAParametersDict(
+            z=self.net, decay=ema_val, device=device
+        )
+
+        # super().__init__(
+        #    width = width * scale,
+        #    height = height * scale,
+        #    tensor = z,
+        #    decay = ema_val,
+        #    device=device,
+        # )
 
     # def get_image_tensor(self):
     def decode_tensor(self):
@@ -129,17 +156,34 @@ def get_latent_tensor(self, detach=False):
         return params
 
     def clone(self):
-        # dummy = super().__init__(*self.image_shape)
+        # dummy = VQGANImage(*self.image_shape)
         # with torch.no_grad():
-        #    #dummy.tensor.set_(self.tensor.clone())
-        #    dummy.net.copy_(self.net)
-        #    dummy.accum.set_(self.accum.clone())
-        #    dummy.biased.set_(self.biased.clone())
-        #    dummy.average.set_(self.average.clone())
-        #    dummy.decay = self.decay
-        dummy = deepcopy(self)
+        #     dummy.representation_parameters.set_(self.representation_parameters.clone())
+        #     dummy.accum.set_(self.accum.clone())
+        #     dummy.biased.set_(self.biased.clone())
+        #     dummy.average.set_(self.average.clone())
+        #     dummy.decay = self.decay
+        # return dummy
+        dummy = DeepImagePrior(*self.image_shape)
+        with torch.no_grad():
+            # dummy.representation_parameters.set_(self.representation_parameters.clone())
+            dummy.image_representation_parameters.set_(
+                self.image_representation_parameters.clone()
+            )
         return dummy
 
+    # def clone(self):
+    #     # dummy = super().__init__(*self.image_shape)
+    #     # with torch.no_grad():
+    #     #    #dummy.tensor.set_(self.tensor.clone())
+    #     #    dummy.net.copy_(self.net)
+    #     #    dummy.accum.set_(self.accum.clone())
+    #     #    dummy.biased.set_(self.biased.clone())
+    #     #    dummy.average.set_(self.average.clone())
+    #     #    dummy.decay = self.decay
+    #     dummy = deepcopy(self)
+    #     return dummy
+
     def encode_random(self):
         pass
 

src/pytti/image_models/ema.py

@@ -58,13 +58,15 @@ def reset(self):
         self.update()
 
 
-class EMAParametersDict(ImageRepresentationalParameters):
+# class EMAParametersDict(ImageRepresentationalParameters):
+class EMAParametersDict(nn.Module):
     """
     LatentTensor with a singleton dimension for the EMAParameters
     """
 
     def __init__(self, z=None, decay=0.99, device=None):
-        super(ImageRepresentationalParameters).__init__()
+        # super(ImageRepresentationalParameters).__init__()
+        super().__init__()
         self.decay = decay
         if device is None:
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -74,9 +76,10 @@ def __init__(self, z=None, decay=0.99, device=None):
     def _new(self, z=None):
         if z is None:
             # I think this can all go in the constructor and doesn't need to call .to()
-            z = torch.zeros(1, 3, self.height, self.width).to(
-                device=self.device, memory_format=torch.channels_last
-            )
+            return nn.Parameter()
+            # z = torch.zeros(1, 3, self.height, self.width).to(
+            #    device=self.device, memory_format=torch.channels_last
+            # )
         # d_ = z
         d_ = {}
         if isinstance(z, EMAParametersDict):
@@ -147,7 +150,7 @@ def reset(self):
 
 
 class EMAImage(DifferentiableImage):
-    def __init__(self, width, height, tensor, decay, device=None):
+    def __init__(self, width, height, tensor=None, decay=0.99, device=None):
         super().__init__(width=width, height=height, device=device)
         self.image_representation_parameters = EMAParametersDict(
             z=tensor, decay=decay, device=device