Update app.py

app.py

@@ -352,7 +352,6 @@ def validate_and_convert_image(image, target_size=(512 , 512)):
 
 class Drag:
 
-    @spaces.GPU
     def __init__(self, device, args, height, width, model_length, dtype=torch.float16, use_sift=False):
         self.device = device
         self.dtype = dtype
@@ -363,12 +362,12 @@ class Drag:
             low_cpu_mem_usage=True,
             custom_resume=True,
         )
-        unet = unet.to(device, dtype)
+        unet = unet.to(dtype)
 
         controlnet = ControlNetSVDModel.from_pretrained(
             os.path.join(args.model, "controlnet"),
         )
-        controlnet = controlnet.to(device, dtype)
+        controlnet = controlnet.to(dtype)
 
         if is_xformers_available():
             import xformers
@@ -386,7 +385,6 @@ class Drag:
             low_cpu_mem_usage=False,
             torch_dtype=torch.float16, variant="fp16", local_files_only=True,
         )
-        pipe.to(device)
 
         self.pipeline = pipe
         # self.pipeline.enable_model_cpu_offload()
@@ -399,6 +397,9 @@ class Drag:
 
     @spaces.GPU
     def run(self, first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id):
+        
+        self.pipeline.to(self.device)
+        
         original_width, original_height = 512, 320  # TODO
 
         # load_image