test

handler.py

@@ -31,29 +31,18 @@ class EndpointHandler():
         #self.pipe.enable_vae_tiling()
         self.generator = torch.Generator(device=device.type).manual_seed(3)
 
-
-        from typing import Optional
-        from torch import Tensor
-        from torch.nn import functional as F
-        from torch.nn import Conv2d
-        from torch.nn.modules.utils import _pair
-
-        def asymmetricConv2DConvForward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):
-            self.paddingX = (self._reversed_padding_repeated_twice[0], self._reversed_padding_repeated_twice[1], 0, 0)
-            self.paddingY = (0, 0, self._reversed_padding_repeated_twice[2], self._reversed_padding_repeated_twice[3])
-            working = F.pad(input, self.paddingX, mode='circular')
-            working = F.pad(working, self.paddingY, mode='constant')
-            return F.conv2d(working, weight, bias, self.stride, _pair(0), self.dilation, self.groups)
-
-        targets = [pipe.vae, pipe.text_encoder, pipe.unet,]
-        conv_layers = []
+        targets = [
+            self.pipe.vae,
+            self.pipe.text_encoder,
+            self.pipe.unet,
+        ]
+        self.conv_layers = []
+        self.conv_layers_original_paddings = []
         for target in targets:
             for module in target.modules():
-                if isinstance(module, torch.nn.Conv2d):
-                    conv_layers.append(module)
-
-        for cl in conv_layers:
-            cl._conv_forward = asymmetricConv2DConvForward.__get__(cl, Conv2d)
+                if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.ConvTranspose2d):
+                    self.conv_layers.append(module)
+                    self.conv_layers_original_paddings.append(module.padding_mode)
 
 
      def __call__(self, data: Any) -> List[List[Dict[str, float]]]: