test

handler.py

@@ -72,11 +72,18 @@ class EndpointHandler():
         self.stable_diffusion_id = "Lykon/dreamshaper-8"
 
         self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id,
-                                                                      #controlnet=self.controlnet, 
+                                                                      controlnet=self.controlnet, 
                                                                       torch_dtype=dtype,
                                                                       safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
         # Define Generator with seed
         self.generator = torch.Generator(device=device.type).manual_seed(3)
+        
+        targets = [self.pipe.vae, self.pipe.unet]
+        for target in targets:
+            for module in target.modules():
+                if isinstance(module, torch.nn.Conv2d):
+                    module.padding_mode = "circular"
+        
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -100,12 +107,7 @@ class EndpointHandler():
             self.pipe.controlnet = self.controlnet
         
         
-        targets = [self.pipe.vae, self.pipe.unet]
-        for target in targets:
-            for module in target.modules():
-                if isinstance(module, torch.nn.Conv2d):
-                    module.padding_mode = "circular"
-        
+
         # hyperparamters
         num_inference_steps = data.pop("num_inference_steps", 30)
         guidance_scale = data.pop("guidance_scale", 7.4)