Update model.py

model.py

@@ -92,7 +92,7 @@ class Model:
         preprocessor_name: str,
     ) -> list[PIL.Image.Image]:
         self.t5_model.model.to('cuda').to(torch.bfloat16)
-        self.gpt_model_canny.to('cuda').to(torch.bfloat16)
+        self.gpt_model_edge.to('cuda').to(torch.bfloat16)
         self.vq_model.to('cuda')
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
@@ -147,6 +147,7 @@ class Model:
             top_k=top_k,
             top_p=top_p,
             sample_logits=True,
+            control_strength=control_strength,
         )
         sampling_time = time.time() - t1
         print(f"Full sampling takes about {sampling_time:.2f} seconds.")
@@ -183,7 +184,7 @@ class Model:
         control_strength: float,
         preprocessor_name: str
     ) -> list[PIL.Image.Image]:
-        self.gpt_model_canny.to('cpu')
+        self.gpt_model_edge.to('cpu')
         self.t5_model.model.to(self.device)
         self.gpt_model_depth.to(self.device)
         self.get_control_depth.model.to(self.device)
@@ -237,6 +238,7 @@ class Model:
             top_k=top_k,
             top_p=top_p,
             sample_logits=True,
+            control_strength=control_strength,
         )
         sampling_time = time.time() - t1
         print(f"Full sampling takes about {sampling_time:.2f} seconds.")