init

app.py

@@ -149,93 +149,31 @@ class Drag:
 
         self.height = height
         self.width = width
-        self.pipeline = None
-        self.cmp = None
-    
-    @spaces.GPU(duration=100)
-    def init_models(self, pretrained_model_name_or_path, resume_from_checkpoint, weight_dtype, device='cuda', enable_xformers_memory_efficient_attention=False, allow_tf32=False):
 
-        print('start loading models...')
-        # Load scheduler, tokenizer and models.
-        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
-            pretrained_model_name_or_path, subfolder="image_encoder", revision=None, variant="fp16"
-        )
-        vae = AutoencoderKLTemporalDecoder.from_pretrained(
-            pretrained_model_name_or_path, subfolder="vae", revision=None, variant="fp16")
-        unet = UNetSpatioTemporalConditionControlNetModel.from_pretrained(
-            pretrained_model_name_or_path,
-            subfolder="unet",
-            low_cpu_mem_usage=True,
-            variant="fp16",
-        )
+    def get_cmp_flow(self, frames, sparse_optical_flow, mask, brush_mask=None):
+
+        '''
+            frames: [b, 13, 3, 384, 384] (0, 1) tensor
+            sparse_optical_flow: [b, 13, 2, 384, 384] (-384, 384) tensor
+            mask: [b, 13, 2, 384, 384] {0, 1} tensor
+        '''
 
-        controlnet = FlowControlNet.from_pretrained(resume_from_checkpoint)
 
         cmp = CMP_demo(
             './models/cmp/experiments/semiauto_annot/resnet50_vip+mpii_liteflow/config.yaml',
             42000
-        ).to(device)
+        ).to('cuda')
         cmp.requires_grad_(False)
 
-        self.cmp = cmp
-        
-        # Freeze vae and image_encoder
-        vae.requires_grad_(False)
-        image_encoder.requires_grad_(False)
-        unet.requires_grad_(False)
-        controlnet.requires_grad_(False)
-
-        # Move image_encoder and vae to gpu and cast to weight_dtype
-        image_encoder.to(device, dtype=weight_dtype)
-        vae.to(device, dtype=weight_dtype)
-        unet.to(device, dtype=weight_dtype)
-        controlnet.to(device, dtype=weight_dtype)
-
-        if enable_xformers_memory_efficient_attention:
-            if is_xformers_available():
-                import xformers
-
-                xformers_version = version.parse(xformers.__version__)
-                if xformers_version == version.parse("0.0.16"):
-                    print(
-                        "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
-                    )
-                unet.enable_xformers_memory_efficient_attention()
-            else:
-                raise ValueError(
-                    "xformers is not available. Make sure it is installed correctly")
-
-        if allow_tf32:
-            torch.backends.cuda.matmul.allow_tf32 = True
-        
-        pipeline = FlowControlNetPipeline.from_pretrained(
-            pretrained_model_name_or_path,
-            unet=unet,
-            controlnet=controlnet,
-            image_encoder=image_encoder,
-            vae=vae,
-            torch_dtype=weight_dtype,
-        )
-        pipeline = pipeline.to(device)
-
-        self.pipeline = pipeline
 
-        print('models loaded.')
 
-    def get_cmp_flow(self, frames, sparse_optical_flow, mask, brush_mask=None):
-
-        '''
-            frames: [b, 13, 3, 384, 384] (0, 1) tensor
-            sparse_optical_flow: [b, 13, 2, 384, 384] (-384, 384) tensor
-            mask: [b, 13, 2, 384, 384] {0, 1} tensor
-        '''
 
         b, t, c, h, w = frames.shape
         assert h == 384 and w == 384
         frames = frames.flatten(0, 1)  # [b*13, 3, 256, 256]
         sparse_optical_flow = sparse_optical_flow.flatten(0, 1)  # [b*13, 2, 256, 256]
         mask = mask.flatten(0, 1)  # [b*13, 2, 256, 256]
-        cmp_flow = self.cmp.run(frames, sparse_optical_flow, mask)  # [b*13, 2, 256, 256]
+        cmp_flow = cmp.run(frames, sparse_optical_flow, mask)  # [b*13, 2, 256, 256]
 
         if brush_mask is not None:
             brush_mask = torch.from_numpy(brush_mask) / 255.
@@ -268,6 +206,54 @@ class Drag:
 
     @torch.no_grad()
     def forward_sample(self, input_drag_384_inmask, input_drag_384_outmask, input_first_frame, input_mask_384_inmask, input_mask_384_outmask, in_mask_flag, out_mask_flag, motion_brush_mask=None, ctrl_scale=1., outputs=dict()):
+
+
+
+
+
+
+
+        pretrained_model_name_or_path="ckpts/stable-video-diffusion-img2vid-xt-1-1"
+        resume_from_checkpoint="ckpts/controlnet"
+        # Load scheduler, tokenizer and models.
+        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+            pretrained_model_name_or_path, subfolder="image_encoder", revision=None, variant="fp16"
+        )
+        vae = AutoencoderKLTemporalDecoder.from_pretrained(
+            pretrained_model_name_or_path, subfolder="vae", revision=None, variant="fp16")
+        unet = UNetSpatioTemporalConditionControlNetModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="unet",
+            low_cpu_mem_usage=True,
+            variant="fp16",
+        )
+        controlnet = FlowControlNet.from_pretrained(resume_from_checkpoint)
+        # Freeze vae and image_encoder
+        vae.requires_grad_(False)
+        image_encoder.requires_grad_(False)
+        unet.requires_grad_(False)
+        controlnet.requires_grad_(False)
+        # Move image_encoder and vae to gpu and cast to weight_dtype
+        image_encoder.to('cuda', dtype=torch.float16)
+        vae.to('cuda', dtype=torch.float16)
+        unet.to('cuda', dtype=torch.float16)
+        controlnet.to('cuda', dtype=torch.float16)
+        # init pipeline
+        pipeline = FlowControlNetPipeline.from_pretrained(
+            pretrained_model_name_or_path,
+            unet=unet,
+            controlnet=controlnet,
+            image_encoder=image_encoder,
+            vae=vae,
+            torch_dtype=torch.float16,
+        )
+        pipeline = pipeline.to('cuda')
+
+
+
+
+
+
         '''
             input_drag: [1, 13, 320, 576, 2]
             input_drag_384: [1, 13, 384, 384, 2]
@@ -321,7 +307,7 @@ class Drag:
 
         controlnet_flow = torch.where(inmask_no_zero, flow_inmask, flow_outmask)
 
-        val_output = self.pipeline(
+        val_output = pipeline(
             input_first_frame_pil, 
             input_first_frame_pil,
             controlnet_flow, 
@@ -358,7 +344,7 @@ class Drag:
 
         return outputs
     
-    @spaces.GPU
+    @spaces.GPU(duration=100)
     @torch.no_grad()
     def get_cmp_flow_from_tracking_points(self, tracking_points, motion_brush_mask, first_frame_path):