Update app.py

app.py

@@ -362,6 +362,7 @@ 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
@@ -379,14 +380,14 @@ class Drag:
         )
         controlnet = controlnet.to(device, dtype)
 
-        # if is_xformers_available():
-        #     import xformers
-        #     xformers_version = version.parse(xformers.__version__)
-        #     unet.enable_xformers_memory_efficient_attention()
-        #     # controlnet.enable_xformers_memory_efficient_attention()
-        # else:
-        #     raise ValueError(
-        #         "xformers is not available. Make sure it is installed correctly")
+        if is_xformers_available():
+            import xformers
+            xformers_version = version.parse(xformers.__version__)
+            unet.enable_xformers_memory_efficient_attention()
+            # controlnet.enable_xformers_memory_efficient_attention()
+        else:
+            raise ValueError(
+                "xformers is not available. Make sure it is installed correctly")
 
         pipe = StableVideoDiffusionInterpControlPipeline.from_pretrained(
             "checkpoints/stable-video-diffusion-img2vid-xt",
@@ -406,6 +407,7 @@ class Drag:
         self.model_length = model_length
         self.use_sift = use_sift
 
+    @spaces.GPU
     def run(self, first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id):
         original_width, original_height = 512, 320  # TODO