update

canonicalize/pipeline_canonicalize.py

@@ -394,8 +394,8 @@ class CanonicalizationPipeline(DiffusionPipeline):
             batch_size = image.shape[0]
         # Define call parameters
         batch_size = 1 if isinstance(prompt, str) else len(prompt)
-        # device = self._execution_device
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        device = self._execution_device
+        # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         import sys
         print(f"PIPELINE Using device!!!!!!!!!!!!: {device}", file=sys.stderr)
         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)

infer_api.py

@@ -875,6 +875,7 @@ infer_canonicalize_validation_pipeline = CanonicalizationPipeline(
 infer_canonicalize_validation_pipeline.set_progress_bar_config(disable=True)
 
 
+@spaces.GPU
 def infer_canonicalize_gen(img_input, seed=0):
     if np.array(img_input).shape[-1] == 4 and np.array(img_input)[..., 3].min() == 255:
         # convert to RGB