init

app.py

@@ -221,12 +221,10 @@ class Drag:
         svd_ckpt = "ckpts/stable-video-diffusion-img2vid-xt-1-1"
         mofa_ckpt = "ckpts/controlnet"
 
-        self.weight_dtype = torch.float16
-
         self.pipeline, self.cmp = init_models(
             svd_ckpt, 
             mofa_ckpt, 
-            weight_dtype=self.weight_dtype, 
+            weight_dtype=torch.float16, 
             device='cuda'
         )
 
@@ -302,12 +300,12 @@ class Drag:
         
         print('start diffusion process...')
 
-        input_drag_384_inmask = input_drag_384_inmask.to('cuda', dtype=self.weight_dtype)
-        mask_384_inmask = mask_384_inmask.to('cuda', dtype=self.weight_dtype)
-        input_drag_384_outmask = input_drag_384_outmask.to('cuda', dtype=self.weight_dtype)
-        mask_384_outmask = mask_384_outmask.to('cuda', dtype=self.weight_dtype)
+        input_drag_384_inmask = input_drag_384_inmask.to('cuda', dtype=torch.float16)
+        mask_384_inmask = mask_384_inmask.to('cuda', dtype=torch.float16)
+        input_drag_384_outmask = input_drag_384_outmask.to('cuda', dtype=torch.float16)
+        mask_384_outmask = mask_384_outmask.to('cuda', dtype=torch.float16)
 
-        input_first_frame_384 = input_first_frame_384.to('cuda', dtype=self.weight_dtype)
+        input_first_frame_384 = input_first_frame_384.to('cuda', dtype=torch.float16)
 
         if in_mask_flag:
             flow_inmask = self.get_flow(
@@ -316,7 +314,7 @@ class Drag:
             )
         else:
             fb, fl = mask_384_inmask.shape[:2]
-            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=self.weight_dtype)
+            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
 
         if out_mask_flag:
             flow_outmask = self.get_flow(
@@ -325,7 +323,7 @@ class Drag:
             )
         else:
             fb, fl = mask_384_outmask.shape[:2]
-            flow_outmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=self.weight_dtype)
+            flow_outmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
         
         inmask_no_zero = (flow_inmask != 0).all(dim=2)
         inmask_no_zero = inmask_no_zero.unsqueeze(2).expand_as(flow_inmask)
@@ -450,12 +448,12 @@ class Drag:
         input_drag_384_outmask = input_drag_384_outmask.permute(0, 1, 4, 2, 3)  # [1, 13, 2, 384, 384]
         mask_384_outmask = input_mask_384_outmask.unsqueeze(2).repeat(1, 1, 2, 1, 1)  # [1, 13, 2, 384, 384]
 
-        input_drag_384_inmask = input_drag_384_inmask.to('cuda', dtype=self.weight_dtype)
-        mask_384_inmask = mask_384_inmask.to('cuda', dtype=self.weight_dtype)
-        input_drag_384_outmask = input_drag_384_outmask.to('cuda', dtype=self.weight_dtype)
-        mask_384_outmask = mask_384_outmask.to('cuda', dtype=self.weight_dtype)
+        input_drag_384_inmask = input_drag_384_inmask.to('cuda', dtype=torch.float16)
+        mask_384_inmask = mask_384_inmask.to('cuda', dtype=torch.float16)
+        input_drag_384_outmask = input_drag_384_outmask.to('cuda', dtype=torch.float16)
+        mask_384_outmask = mask_384_outmask.to('cuda', dtype=torch.float16)
 
-        input_first_frame_384 = input_first_frame_384.to('cuda', dtype=self.weight_dtype)
+        input_first_frame_384 = input_first_frame_384.to('cuda', dtype=torch.float16)
 
         if in_mask_flag:
             flow_inmask = self.get_flow(
@@ -464,7 +462,7 @@ class Drag:
             )
         else:
             fb, fl = mask_384_inmask.shape[:2]
-            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=self.weight_dtype)
+            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
 
         if out_mask_flag:
             flow_outmask = self.get_flow(
@@ -473,7 +471,7 @@ class Drag:
             )
         else:
             fb, fl = mask_384_outmask.shape[:2]
-            flow_outmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=self.weight_dtype)
+            flow_outmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
         
         inmask_no_zero = (flow_inmask != 0).all(dim=2)
         inmask_no_zero = inmask_no_zero.unsqueeze(2).expand_as(flow_inmask)