Update app.py

app.py

@@ -352,6 +352,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
@@ -362,21 +363,21 @@ class Drag:
             low_cpu_mem_usage=True,
             custom_resume=True,
         )
-        unet = unet.to(dtype)
+        unet = unet.to(device, dtype)
 
         controlnet = ControlNetSVDModel.from_pretrained(
             os.path.join(args.model, "controlnet"),
         )
-        controlnet = controlnet.to(dtype)
+        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",
@@ -385,6 +386,7 @@ 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()
@@ -396,10 +398,7 @@ class Drag:
         self.use_sift = use_sift
 
     @spaces.GPU
-    def run(self, first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id):
-        
-        self.pipeline.to(self.device)
-        
+    def run(self, first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id):        
         original_width, original_height = 512, 320  # TODO
 
         # load_image
@@ -530,7 +529,7 @@ class Drag:
 def reset_states(first_frame_path, last_frame_path, tracking_points):
     first_frame_path = gr.State()
     last_frame_path = gr.State()
-    tracking_points = gr.State()
+    tracking_points = gr.State([])
 
     return first_frame_path, last_frame_path, tracking_points
 
@@ -549,7 +548,7 @@ def preprocess_image(image):
     
     image_pil.save(first_frame_path)
 
-    return first_frame_path, first_frame_path, gr.State()
+    return first_frame_path, first_frame_path, gr.State([])
 
 
 def preprocess_image_end(image_end):
@@ -566,7 +565,7 @@ def preprocess_image_end(image_end):
 
     image_end_pil.save(last_frame_path)
 
-    return last_frame_path, last_frame_path, gr.State()
+    return last_frame_path, last_frame_path, gr.State([])
 
 
 def add_drag(tracking_points):
@@ -680,7 +679,6 @@ if __name__ == "__main__":
     args = get_args()
     ensure_dirname(args.output_dir)
     
-    
     color_list = []
     for i in range(20):
         color = np.concatenate([np.random.random(4)*255], axis=0)
@@ -710,7 +708,7 @@ if __name__ == "__main__":
         Framer = Drag("cuda", args, 320, 512, 14)
         first_frame_path = gr.State()
         last_frame_path = gr.State()
-        tracking_points = gr.State()
+        tracking_points = gr.State([])
     
         with gr.Row():
             with gr.Column(scale=1):