minor optimization

server/pipelines/pix2pixTurbo.py

@@ -7,7 +7,6 @@ from PIL import Image
 from pipelines.pix2pix.pix2pix_turbo import Pix2Pix_Turbo
 from pipelines.utils.canny_gpu import SobelOperator
 
-
 default_prompt = "close-up photo of the joker"
 page_content = """
 <h1 class="text-3xl font-bold">Real-Time pix2pix_turbo</h1>
@@ -94,6 +93,7 @@ class Pipeline:
         self.model = Pix2Pix_Turbo("edge_to_image")
         self.canny_torch = SobelOperator(device=device)
         self.device = device
+        self.last_time = 0.0
 
     def predict(self, params: "Pipeline.InputParams") -> Image.Image:
         # generator = torch.manual_seed(params.seed)
@@ -123,5 +123,4 @@ class Pipeline:
             control_image = canny_pil.resize((w0, h0))
             w1, h1 = result_image.size
             result_image.paste(control_image, (w1 - w0, h1 - h0))
-
         return result_image

server/pipelines/utils/canny_gpu.py

@@ -5,6 +5,13 @@ from PIL import Image
 
 
 class SobelOperator(nn.Module):
+    SOBEL_KERNEL_X = torch.tensor(
+        [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]]
+    )
+    SOBEL_KERNEL_Y = torch.tensor(
+        [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]]
+    )
+
     def __init__(self, device="cuda"):
         super(SobelOperator, self).__init__()
         self.device = device
@@ -14,17 +21,13 @@ class SobelOperator(nn.Module):
         self.edge_conv_y = nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False).to(
             self.device
         )
-
-        sobel_kernel_x = torch.tensor(
-            [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], device=self.device
+        self.edge_conv_x.weight = nn.Parameter(
+            self.SOBEL_KERNEL_X.view((1, 1, 3, 3)).to(self.device)
         )
-        sobel_kernel_y = torch.tensor(
-            [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]], device=self.device
+        self.edge_conv_y.weight = nn.Parameter(
+            self.SOBEL_KERNEL_Y.view((1, 1, 3, 3)).to(self.device)
         )
 
-        self.edge_conv_x.weight = nn.Parameter(sobel_kernel_x.view((1, 1, 3, 3)))
-        self.edge_conv_y.weight = nn.Parameter(sobel_kernel_y.view((1, 1, 3, 3)))
-
     @torch.no_grad()
     def forward(
         self,
@@ -40,10 +43,10 @@ class SobelOperator(nn.Module):
         # Compute gradients
         edge_x = self.edge_conv_x(image_tensor)
         edge_y = self.edge_conv_y(image_tensor)
-        edge = torch.sqrt(edge_x**2 + edge_y**2)
+        edge = torch.sqrt(torch.square(edge_x) + torch.square(edge_y))
 
         # Apply thresholding
-        edge = edge / edge.max()  # Normalize to 0-1
+        edge.div_(edge.max())  # Normalize to 0-1 (in-place operation)
         edge[edge >= high_threshold] = 1.0
         edge[edge <= low_threshold] = 0.0