init

app.py

@@ -89,6 +89,79 @@ def get_sparseflow_and_mask_forward(
     return s_flow, mask
 
 
+@spaces.GPU(duration=100)
+def init_models(pretrained_model_name_or_path, resume_from_checkpoint, weight_dtype, device='cuda', enable_xformers_memory_efficient_attention=False, allow_tf32=False):
+
+    from models.unet_spatio_temporal_condition_controlnet import UNetSpatioTemporalConditionControlNetModel
+    from pipeline.pipeline import FlowControlNetPipeline
+    from models.svdxt_featureflow_forward_controlnet_s2d_fixcmp_norefine import FlowControlNet, CMP_demo
+
+    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",
+    )
+
+    controlnet = FlowControlNet.from_pretrained(resume_from_checkpoint)
+
+    cmp = CMP_demo(
+        './models/cmp/experiments/semiauto_annot/resnet50_vip+mpii_liteflow/config.yaml',
+        42000
+    ).to(device)
+    cmp.requires_grad_(False)
+    
+    # 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)
+
+    print('models loaded.')
+
+    return pipeline, cmp
+
+
 def interpolate_trajectory(points, n_points):
     x = [point[0] for point in points]
     y = [point[1] for point in points]
@@ -142,110 +215,22 @@ def visualize_drag_v2(background_image_path, splited_tracks, width, height):
     return trajectory_maps, transparent_layer
 
 
-with gr.Blocks() as demo:
-    gr.Markdown("""<h1 align="center">MOFA-Video</h1><br>""")
-
-    gr.Markdown("""Official Gradio Demo for <a href='https://myniuuu.github.io/MOFA_Video'><b>MOFA-Video: Controllable Image Animation via Generative Motion Field Adaptions in Frozen Image-to-Video Diffusion Model</b></a>.<br>""")
-
-    gr.Markdown(
-        """
-        During the inference, kindly follow these instructions:
-        <br>
-        1. Use the "Upload Image" button to upload an image. Avoid dragging the image directly into the window. <br>
-        2. Proceed to draw trajectories: <br>
-            2.1. Click "Add Trajectory" first, then select points on the "Add Trajectory Here" image. The first click sets the starting point. Click multiple points to create a non-linear trajectory. To add a new trajectory, click "Add Trajectory" again and select points on the image. Avoid clicking the "Add Trajectory" button multiple times without clicking points in the image to add the trajectory, as this can lead to errors. <br>
-            2.2. After adding each trajectory, an optical flow image will be displayed automatically. Use it as a reference to adjust the trajectory for desired effects (e.g., area, intensity). <br>
-            2.3. To delete the latest trajectory, click "Delete Last Trajectory." <br>
-            2.4. Choose the Control Scale in the bar. This determines the control intensity. Setting it to 0 means no control (pure generation result of SVD itself), while setting it to 1 results in the strongest control (which will not lead to good results in most cases because of twisting artifacts). A preset value of 0.6 is recommended for most cases. <br>
-            2.5. To use the motion brush for restraining the control area of the trajectory, click to add masks on the "Add Motion Brush Here" image. The motion brush restricts the optical flow area derived from the trajectory whose starting point is within the motion brush. The displayed optical flow image will change correspondingly. Adjust the motion brush radius using the "Motion Brush Radius" bar. <br>
-        3. Click the "Run" button to animate the image according to the path. <br>
-        """
-    )
 
-    height, width = 512, 512
+pipeline, cmp = init_models(
+    "ckpts/stable-video-diffusion-img2vid-xt-1-1", 
+    "ckpts/controlnet", 
+    weight_dtype=torch.float16, 
+    device='cuda'
+)
 
-    pipeline, cmp = None, None
 
-    first_frame_path = gr.State()
-    tracking_points = gr.State([])
-    motion_brush_points = gr.State([])
-    motion_brush_mask = gr.State()
-    motion_brush_viz = gr.State()
-    inference_batch_size = gr.State(1)
+class Drag:
+    def __init__(self, height, width):
 
-    @spaces.GPU(duration=100)
-    def init_models(pretrained_model_name_or_path="ckpts/stable-video-diffusion-img2vid-xt-1-1", resume_from_checkpoint="ckpts/controlnet", weight_dtype=torch.float16, device='cuda', enable_xformers_memory_efficient_attention=False, allow_tf32=False):
+        self.height = height
+        self.width = width
 
-        from models.unet_spatio_temporal_condition_controlnet import UNetSpatioTemporalConditionControlNetModel
-        from pipeline.pipeline import FlowControlNetPipeline
-        from models.svdxt_featureflow_forward_controlnet_s2d_fixcmp_norefine import FlowControlNet, CMP_demo
-
-        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",
-        )
-
-        controlnet = FlowControlNet.from_pretrained(resume_from_checkpoint)
-
-        cmp = CMP_demo(
-            './models/cmp/experiments/semiauto_annot/resnet50_vip+mpii_liteflow/config.yaml',
-            42000
-        ).to(device)
-        cmp.requires_grad_(False)
-        
-        # 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)
-
-        print('models loaded.')
-
-        return pipeline, cmp
-
-    def get_cmp_flow(frames, sparse_optical_flow, mask, brush_mask=None):
+    def get_cmp_flow(self, frames, sparse_optical_flow, mask, brush_mask=None):
 
         '''
             frames: [b, 13, 3, 384, 384] (0, 1) tensor
@@ -270,19 +255,19 @@ with gr.Blocks() as demo:
         return cmp_flow
     
 
-    def get_flow(pixel_values_384, sparse_optical_flow_384, mask_384, motion_brush_mask=None):
+    def get_flow(self, pixel_values_384, sparse_optical_flow_384, mask_384, motion_brush_mask=None):
 
         fb, fl, fc, _, _ = pixel_values_384.shape
 
-        controlnet_flow = get_cmp_flow(
+        controlnet_flow = self.get_cmp_flow(
             pixel_values_384[:, 0:1, :, :, :].repeat(1, fl, 1, 1, 1), 
             sparse_optical_flow_384, 
             mask_384, motion_brush_mask
         )
 
-        if height != 384 or width != 384:
-            scales = [height / 384, width / 384]
-            controlnet_flow = F.interpolate(controlnet_flow.flatten(0, 1), (height, width), mode='nearest').reshape(fb, fl, 2, height, width)
+        if self.height != 384 or self.width != 384:
+            scales = [self.height / 384, self.width / 384]
+            controlnet_flow = F.interpolate(controlnet_flow.flatten(0, 1), (self.height, self.width), mode='nearest').reshape(fb, fl, 2, self.height, self.width)
             controlnet_flow[:, :, 0] *= scales[1]
             controlnet_flow[:, :, 1] *= scales[0]
         
@@ -290,7 +275,7 @@ with gr.Blocks() as demo:
     
 
     @torch.no_grad()
-    def forward_sample(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()):
+    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()):
         '''
             input_drag: [1, 13, 320, 576, 2]
             input_drag_384: [1, 13, 384, 384, 2]
@@ -322,22 +307,22 @@ with gr.Blocks() as demo:
         input_first_frame_384 = input_first_frame_384.to('cuda', dtype=torch.float16)
 
         if in_mask_flag:
-            flow_inmask = get_flow(
+            flow_inmask = self.get_flow(
                 input_first_frame_384, 
                 input_drag_384_inmask, mask_384_inmask, motion_brush_mask
             )
         else:
             fb, fl = mask_384_inmask.shape[:2]
-            flow_inmask = torch.zeros(fb, fl, 2, height, width).to('cuda', dtype=torch.float16)
+            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
 
         if out_mask_flag:
-            flow_outmask = get_flow(
+            flow_outmask = self.get_flow(
                 input_first_frame_384, 
                 input_drag_384_outmask, mask_384_outmask
             )
         else:
             fb, fl = mask_384_outmask.shape[:2]
-            flow_outmask = torch.zeros(fb, fl, 2, height, width).to('cuda', dtype=torch.float16)
+            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)
@@ -383,16 +368,16 @@ with gr.Blocks() as demo:
     
     @spaces.GPU
     @torch.no_grad()
-    def get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path):
+    def get_cmp_flow_from_tracking_points(self, tracking_points, motion_brush_mask, first_frame_path):
 
-        original_width, original_height = width, height
+        original_width, original_height = self.width, self.height
 
         input_all_points = tracking_points.constructor_args['value']
 
         if len(input_all_points) == 0 or len(input_all_points[-1]) == 1:
             return np.uint8(np.ones((original_width, original_height, 3))*255)
         
-        resized_all_points = [tuple([tuple([int(e1[0]*width/original_width), int(e1[1]*height/original_height)]) for e1 in e]) for e in input_all_points]
+        resized_all_points = [tuple([tuple([int(e1[0]*self.width/original_width), int(e1[1]*self.height/original_height)]) for e1 in e]) for e in input_all_points]
         resized_all_points_384 = [tuple([tuple([int(e1[0]*384/original_width), int(e1[1]*384/original_height)]) for e1 in e]) for e in input_all_points]
 
         new_resized_all_points = []
@@ -470,22 +455,22 @@ with gr.Blocks() as demo:
         input_first_frame_384 = input_first_frame_384.to('cuda', dtype=torch.float16)
 
         if in_mask_flag:
-            flow_inmask = get_flow(
+            flow_inmask = self.get_flow(
                 input_first_frame_384, 
                 input_drag_384_inmask, mask_384_inmask, motion_brush_mask_384
             )
         else:
             fb, fl = mask_384_inmask.shape[:2]
-            flow_inmask = torch.zeros(fb, fl, 2, height, width).to('cuda', dtype=torch.float16)
+            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
 
         if out_mask_flag:
-            flow_outmask = get_flow(
+            flow_outmask = self.get_flow(
                 input_first_frame_384, 
                 input_drag_384_outmask, mask_384_outmask
             )
         else:
             fb, fl = mask_384_outmask.shape[:2]
-            flow_outmask = torch.zeros(fb, fl, 2, height, width).to('cuda', dtype=torch.float16)
+            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)
@@ -498,12 +483,12 @@ with gr.Blocks() as demo:
         return viz_esti_flows
     
     @spaces.GPU(duration=200)
-    def run(first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale):
+    def run(self, first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale):
         
-        original_width, original_height = width, height
+        original_width, original_height = self.width, self.height
 
         input_all_points = tracking_points.constructor_args['value']
-        resized_all_points = [tuple([tuple([int(e1[0]*width/original_width), int(e1[1]*height/original_height)]) for e1 in e]) for e in input_all_points]
+        resized_all_points = [tuple([tuple([int(e1[0]*self.width/original_width), int(e1[1]*self.height/original_height)]) for e1 in e]) for e in input_all_points]
         resized_all_points_384 = [tuple([tuple([int(e1[0]*384/original_width), int(e1[1]*384/original_height)]) for e1 in e]) for e in input_all_points]
 
         new_resized_all_points = []
@@ -556,9 +541,9 @@ with gr.Blocks() as demo:
         id = base.split('_')[0]
         
         image_pil = image2pil(first_frame_path)
-        image_pil = image_pil.resize((width, height), Image.BILINEAR).convert('RGB')
+        image_pil = image_pil.resize((self.width, self.height), Image.BILINEAR).convert('RGB')
         
-        visualized_drag, _ = visualize_drag_v2(first_frame_path, resized_all_points, width, height)
+        visualized_drag, _ = visualize_drag_v2(first_frame_path, resized_all_points, self.width, self.height)
 
         motion_brush_viz_pil = Image.fromarray(motion_brush_viz.astype(np.uint8)).convert('RGBA')
         visualized_drag = visualized_drag[0].convert('RGBA')
@@ -581,7 +566,7 @@ with gr.Blocks() as demo:
                 first_frames = outputs['logits_imgs'][:, -1]
             
 
-            outputs = forward_sample(
+            outputs = self.forward_sample(
                 input_drag_384_inmask.to('cuda'), 
                 input_drag_384_outmask.to('cuda'), 
                 first_frames.to('cuda'),
@@ -644,16 +629,43 @@ with gr.Blocks() as demo:
 
         return hint_path, outputs_path, flows_path, outputs_mp4_path, flows_mp4_path
 
-    @spaces.GPU(duration=100)
-    def preprocess_image(image):
 
-        pipeline, cmp = init_models()
+with gr.Blocks() as demo:
+    gr.Markdown("""<h1 align="center">MOFA-Video</h1><br>""")
+
+    gr.Markdown("""Official Gradio Demo for <a href='https://myniuuu.github.io/MOFA_Video'><b>MOFA-Video: Controllable Image Animation via Generative Motion Field Adaptions in Frozen Image-to-Video Diffusion Model</b></a>.<br>""")
+
+    gr.Markdown(
+        """
+        During the inference, kindly follow these instructions:
+        <br>
+        1. Use the "Upload Image" button to upload an image. Avoid dragging the image directly into the window. <br>
+        2. Proceed to draw trajectories: <br>
+            2.1. Click "Add Trajectory" first, then select points on the "Add Trajectory Here" image. The first click sets the starting point. Click multiple points to create a non-linear trajectory. To add a new trajectory, click "Add Trajectory" again and select points on the image. Avoid clicking the "Add Trajectory" button multiple times without clicking points in the image to add the trajectory, as this can lead to errors. <br>
+            2.2. After adding each trajectory, an optical flow image will be displayed automatically. Use it as a reference to adjust the trajectory for desired effects (e.g., area, intensity). <br>
+            2.3. To delete the latest trajectory, click "Delete Last Trajectory." <br>
+            2.4. Choose the Control Scale in the bar. This determines the control intensity. Setting it to 0 means no control (pure generation result of SVD itself), while setting it to 1 results in the strongest control (which will not lead to good results in most cases because of twisting artifacts). A preset value of 0.6 is recommended for most cases. <br>
+            2.5. To use the motion brush for restraining the control area of the trajectory, click to add masks on the "Add Motion Brush Here" image. The motion brush restricts the optical flow area derived from the trajectory whose starting point is within the motion brush. The displayed optical flow image will change correspondingly. Adjust the motion brush radius using the "Motion Brush Radius" bar. <br>
+        3. Click the "Run" button to animate the image according to the path. <br>
+        """
+    )
+
+    target_size = 512
+    DragNUWA_net = Drag(target_size, target_size)
+    first_frame_path = gr.State()
+    tracking_points = gr.State([])
+    motion_brush_points = gr.State([])
+    motion_brush_mask = gr.State()
+    motion_brush_viz = gr.State()
+    inference_batch_size = gr.State(1)
+
+    def preprocess_image(image):
 
         image_pil = image2pil(image.name)
         raw_w, raw_h = image_pil.size
 
         max_edge = min(raw_w, raw_h)
-        resize_ratio = width / max_edge
+        resize_ratio = target_size / max_edge
 
         image_pil = image_pil.resize((round(raw_w * resize_ratio), round(raw_h * resize_ratio)), Image.BILINEAR)
 
@@ -663,8 +675,8 @@ with gr.Blocks() as demo:
 
         image_pil = transforms.CenterCrop((crop_h, crop_w))(image_pil.convert('RGB'))
 
-        width = crop_w
-        height = crop_h
+        DragNUWA_net.width = crop_w
+        DragNUWA_net.height = crop_h
 
         id = str(time.time()).split('.')[0]
         os.makedirs(os.path.join(output_dir_video, str(id)), exist_ok=True)
@@ -709,7 +721,7 @@ with gr.Blocks() as demo:
         transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
         trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
 
-        viz_flow = get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
+        viz_flow = DragNUWA_net.get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
 
         return tracking_points, trajectory_map, viz_flow
     
@@ -729,7 +741,7 @@ with gr.Blocks() as demo:
         transparent_layer_pil = Image.fromarray(transparent_layer.astype(np.uint8))
         motion_map = Image.alpha_composite(transparent_background, transparent_layer_pil)
 
-        viz_flow = get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
+        viz_flow = DragNUWA_net.get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
 
         return motion_brush_mask, transparent_layer, motion_map, viz_flow
 
@@ -765,7 +777,7 @@ with gr.Blocks() as demo:
         transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
         trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
 
-        viz_flow = get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
+        viz_flow = DragNUWA_net.get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
 
         return tracking_points, trajectory_map, viz_flow
 
@@ -820,6 +832,6 @@ with gr.Blocks() as demo:
 
     input_image_mask.select(add_motion_brushes, [motion_brush_points, motion_brush_mask, motion_brush_viz, first_frame_path, brush_radius, tracking_points], [motion_brush_mask, motion_brush_viz, input_image_mask, viz_flow])
 
-    run_button.click(run, [first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale], [hint_image, output_video, output_flow, output_video_mp4, output_flow_mp4])
+    run_button.click(DragNUWA_net.run, [first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale], [hint_image, output_video, output_flow, output_video_mp4, output_flow_mp4])
 
     demo.launch()

oldapp.py → modifiedapp.py

@@ -89,78 +89,6 @@ def get_sparseflow_and_mask_forward(
     return s_flow, mask
 
 
-def init_models(pretrained_model_name_or_path, resume_from_checkpoint, weight_dtype, device='cuda', enable_xformers_memory_efficient_attention=False, allow_tf32=False):
-
-    from models.unet_spatio_temporal_condition_controlnet import UNetSpatioTemporalConditionControlNetModel
-    from pipeline.pipeline import FlowControlNetPipeline
-    from models.svdxt_featureflow_forward_controlnet_s2d_fixcmp_norefine import FlowControlNet, CMP_demo
-
-    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",
-    )
-
-    controlnet = FlowControlNet.from_pretrained(resume_from_checkpoint)
-
-    cmp = CMP_demo(
-        './models/cmp/experiments/semiauto_annot/resnet50_vip+mpii_liteflow/config.yaml',
-        42000
-    ).to(device)
-    cmp.requires_grad_(False)
-    
-    # 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)
-
-    print('models loaded.')
-
-    return pipeline, cmp
-
-
 def interpolate_trajectory(points, n_points):
     x = [point[0] for point in points]
     y = [point[1] for point in points]
@@ -214,24 +142,110 @@ def visualize_drag_v2(background_image_path, splited_tracks, width, height):
     return trajectory_maps, transparent_layer
 
 
-class Drag:
-    @spaces.GPU(duration=200)
-    def __init__(self, height, width):
+with gr.Blocks() as demo:
+    gr.Markdown("""<h1 align="center">MOFA-Video</h1><br>""")
+
+    gr.Markdown("""Official Gradio Demo for <a href='https://myniuuu.github.io/MOFA_Video'><b>MOFA-Video: Controllable Image Animation via Generative Motion Field Adaptions in Frozen Image-to-Video Diffusion Model</b></a>.<br>""")
+
+    gr.Markdown(
+        """
+        During the inference, kindly follow these instructions:
+        <br>
+        1. Use the "Upload Image" button to upload an image. Avoid dragging the image directly into the window. <br>
+        2. Proceed to draw trajectories: <br>
+            2.1. Click "Add Trajectory" first, then select points on the "Add Trajectory Here" image. The first click sets the starting point. Click multiple points to create a non-linear trajectory. To add a new trajectory, click "Add Trajectory" again and select points on the image. Avoid clicking the "Add Trajectory" button multiple times without clicking points in the image to add the trajectory, as this can lead to errors. <br>
+            2.2. After adding each trajectory, an optical flow image will be displayed automatically. Use it as a reference to adjust the trajectory for desired effects (e.g., area, intensity). <br>
+            2.3. To delete the latest trajectory, click "Delete Last Trajectory." <br>
+            2.4. Choose the Control Scale in the bar. This determines the control intensity. Setting it to 0 means no control (pure generation result of SVD itself), while setting it to 1 results in the strongest control (which will not lead to good results in most cases because of twisting artifacts). A preset value of 0.6 is recommended for most cases. <br>
+            2.5. To use the motion brush for restraining the control area of the trajectory, click to add masks on the "Add Motion Brush Here" image. The motion brush restricts the optical flow area derived from the trajectory whose starting point is within the motion brush. The displayed optical flow image will change correspondingly. Adjust the motion brush radius using the "Motion Brush Radius" bar. <br>
+        3. Click the "Run" button to animate the image according to the path. <br>
+        """
+    )
+
+    height, width = 512, 512
+
+    pipeline, cmp = None, None
+
+    first_frame_path = gr.State()
+    tracking_points = gr.State([])
+    motion_brush_points = gr.State([])
+    motion_brush_mask = gr.State()
+    motion_brush_viz = gr.State()
+    inference_batch_size = gr.State(1)
 
-        svd_ckpt = "ckpts/stable-video-diffusion-img2vid-xt-1-1"
-        mofa_ckpt = "ckpts/controlnet"
+    @spaces.GPU(duration=100)
+    def init_models(pretrained_model_name_or_path="ckpts/stable-video-diffusion-img2vid-xt-1-1", resume_from_checkpoint="ckpts/controlnet", weight_dtype=torch.float16, device='cuda', enable_xformers_memory_efficient_attention=False, allow_tf32=False):
 
-        self.pipeline, self.cmp = init_models(
-            svd_ckpt, 
-            mofa_ckpt, 
-            weight_dtype=torch.float16, 
-            device='cuda'
+        from models.unet_spatio_temporal_condition_controlnet import UNetSpatioTemporalConditionControlNetModel
+        from pipeline.pipeline import FlowControlNetPipeline
+        from models.svdxt_featureflow_forward_controlnet_s2d_fixcmp_norefine import FlowControlNet, CMP_demo
+
+        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",
         )
 
-        self.height = height
-        self.width = width
+        controlnet = FlowControlNet.from_pretrained(resume_from_checkpoint)
+
+        cmp = CMP_demo(
+            './models/cmp/experiments/semiauto_annot/resnet50_vip+mpii_liteflow/config.yaml',
+            42000
+        ).to(device)
+        cmp.requires_grad_(False)
+        
+        # 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)
 
-    def get_cmp_flow(self, frames, sparse_optical_flow, mask, brush_mask=None):
+        print('models loaded.')
+
+        return pipeline, cmp
+
+    def get_cmp_flow(frames, sparse_optical_flow, mask, brush_mask=None):
 
         '''
             frames: [b, 13, 3, 384, 384] (0, 1) tensor
@@ -244,7 +258,7 @@ class Drag:
         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.
@@ -256,19 +270,19 @@ class Drag:
         return cmp_flow
     
 
-    def get_flow(self, pixel_values_384, sparse_optical_flow_384, mask_384, motion_brush_mask=None):
+    def get_flow(pixel_values_384, sparse_optical_flow_384, mask_384, motion_brush_mask=None):
 
         fb, fl, fc, _, _ = pixel_values_384.shape
 
-        controlnet_flow = self.get_cmp_flow(
+        controlnet_flow = get_cmp_flow(
             pixel_values_384[:, 0:1, :, :, :].repeat(1, fl, 1, 1, 1), 
             sparse_optical_flow_384, 
             mask_384, motion_brush_mask
         )
 
-        if self.height != 384 or self.width != 384:
-            scales = [self.height / 384, self.width / 384]
-            controlnet_flow = F.interpolate(controlnet_flow.flatten(0, 1), (self.height, self.width), mode='nearest').reshape(fb, fl, 2, self.height, self.width)
+        if height != 384 or width != 384:
+            scales = [height / 384, width / 384]
+            controlnet_flow = F.interpolate(controlnet_flow.flatten(0, 1), (height, width), mode='nearest').reshape(fb, fl, 2, height, width)
             controlnet_flow[:, :, 0] *= scales[1]
             controlnet_flow[:, :, 1] *= scales[0]
         
@@ -276,7 +290,7 @@ 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()):
+    def forward_sample(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()):
         '''
             input_drag: [1, 13, 320, 576, 2]
             input_drag_384: [1, 13, 384, 384, 2]
@@ -308,29 +322,29 @@ class Drag:
         input_first_frame_384 = input_first_frame_384.to('cuda', dtype=torch.float16)
 
         if in_mask_flag:
-            flow_inmask = self.get_flow(
+            flow_inmask = get_flow(
                 input_first_frame_384, 
                 input_drag_384_inmask, mask_384_inmask, motion_brush_mask
             )
         else:
             fb, fl = mask_384_inmask.shape[:2]
-            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
+            flow_inmask = torch.zeros(fb, fl, 2, height, width).to('cuda', dtype=torch.float16)
 
         if out_mask_flag:
-            flow_outmask = self.get_flow(
+            flow_outmask = get_flow(
                 input_first_frame_384, 
                 input_drag_384_outmask, mask_384_outmask
             )
         else:
             fb, fl = mask_384_outmask.shape[:2]
-            flow_outmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
+            flow_outmask = torch.zeros(fb, fl, 2, height, 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)
 
         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, 
@@ -369,16 +383,16 @@ class Drag:
     
     @spaces.GPU
     @torch.no_grad()
-    def get_cmp_flow_from_tracking_points(self, tracking_points, motion_brush_mask, first_frame_path):
+    def get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path):
 
-        original_width, original_height = self.width, self.height
+        original_width, original_height = width, height
 
         input_all_points = tracking_points.constructor_args['value']
 
         if len(input_all_points) == 0 or len(input_all_points[-1]) == 1:
             return np.uint8(np.ones((original_width, original_height, 3))*255)
         
-        resized_all_points = [tuple([tuple([int(e1[0]*self.width/original_width), int(e1[1]*self.height/original_height)]) for e1 in e]) for e in input_all_points]
+        resized_all_points = [tuple([tuple([int(e1[0]*width/original_width), int(e1[1]*height/original_height)]) for e1 in e]) for e in input_all_points]
         resized_all_points_384 = [tuple([tuple([int(e1[0]*384/original_width), int(e1[1]*384/original_height)]) for e1 in e]) for e in input_all_points]
 
         new_resized_all_points = []
@@ -456,22 +470,22 @@ class Drag:
         input_first_frame_384 = input_first_frame_384.to('cuda', dtype=torch.float16)
 
         if in_mask_flag:
-            flow_inmask = self.get_flow(
+            flow_inmask = get_flow(
                 input_first_frame_384, 
                 input_drag_384_inmask, mask_384_inmask, motion_brush_mask_384
             )
         else:
             fb, fl = mask_384_inmask.shape[:2]
-            flow_inmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
+            flow_inmask = torch.zeros(fb, fl, 2, height, width).to('cuda', dtype=torch.float16)
 
         if out_mask_flag:
-            flow_outmask = self.get_flow(
+            flow_outmask = get_flow(
                 input_first_frame_384, 
                 input_drag_384_outmask, mask_384_outmask
             )
         else:
             fb, fl = mask_384_outmask.shape[:2]
-            flow_outmask = torch.zeros(fb, fl, 2, self.height, self.width).to('cuda', dtype=torch.float16)
+            flow_outmask = torch.zeros(fb, fl, 2, height, 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)
@@ -484,12 +498,12 @@ class Drag:
         return viz_esti_flows
     
     @spaces.GPU(duration=200)
-    def run(self, first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale):
+    def run(first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale):
         
-        original_width, original_height = self.width, self.height
+        original_width, original_height = width, height
 
         input_all_points = tracking_points.constructor_args['value']
-        resized_all_points = [tuple([tuple([int(e1[0]*self.width/original_width), int(e1[1]*self.height/original_height)]) for e1 in e]) for e in input_all_points]
+        resized_all_points = [tuple([tuple([int(e1[0]*width/original_width), int(e1[1]*height/original_height)]) for e1 in e]) for e in input_all_points]
         resized_all_points_384 = [tuple([tuple([int(e1[0]*384/original_width), int(e1[1]*384/original_height)]) for e1 in e]) for e in input_all_points]
 
         new_resized_all_points = []
@@ -542,9 +556,9 @@ class Drag:
         id = base.split('_')[0]
         
         image_pil = image2pil(first_frame_path)
-        image_pil = image_pil.resize((self.width, self.height), Image.BILINEAR).convert('RGB')
+        image_pil = image_pil.resize((width, height), Image.BILINEAR).convert('RGB')
         
-        visualized_drag, _ = visualize_drag_v2(first_frame_path, resized_all_points, self.width, self.height)
+        visualized_drag, _ = visualize_drag_v2(first_frame_path, resized_all_points, width, height)
 
         motion_brush_viz_pil = Image.fromarray(motion_brush_viz.astype(np.uint8)).convert('RGBA')
         visualized_drag = visualized_drag[0].convert('RGBA')
@@ -567,7 +581,7 @@ class Drag:
                 first_frames = outputs['logits_imgs'][:, -1]
             
 
-            outputs = self.forward_sample(
+            outputs = forward_sample(
                 input_drag_384_inmask.to('cuda'), 
                 input_drag_384_outmask.to('cuda'), 
                 first_frames.to('cuda'),
@@ -630,43 +644,17 @@ class Drag:
 
         return hint_path, outputs_path, flows_path, outputs_mp4_path, flows_mp4_path
 
-
-with gr.Blocks() as demo:
-    gr.Markdown("""<h1 align="center">MOFA-Video</h1><br>""")
-
-    gr.Markdown("""Official Gradio Demo for <a href='https://myniuuu.github.io/MOFA_Video'><b>MOFA-Video: Controllable Image Animation via Generative Motion Field Adaptions in Frozen Image-to-Video Diffusion Model</b></a>.<br>""")
-
-    gr.Markdown(
-        """
-        During the inference, kindly follow these instructions:
-        <br>
-        1. Use the "Upload Image" button to upload an image. Avoid dragging the image directly into the window. <br>
-        2. Proceed to draw trajectories: <br>
-            2.1. Click "Add Trajectory" first, then select points on the "Add Trajectory Here" image. The first click sets the starting point. Click multiple points to create a non-linear trajectory. To add a new trajectory, click "Add Trajectory" again and select points on the image. Avoid clicking the "Add Trajectory" button multiple times without clicking points in the image to add the trajectory, as this can lead to errors. <br>
-            2.2. After adding each trajectory, an optical flow image will be displayed automatically. Use it as a reference to adjust the trajectory for desired effects (e.g., area, intensity). <br>
-            2.3. To delete the latest trajectory, click "Delete Last Trajectory." <br>
-            2.4. Choose the Control Scale in the bar. This determines the control intensity. Setting it to 0 means no control (pure generation result of SVD itself), while setting it to 1 results in the strongest control (which will not lead to good results in most cases because of twisting artifacts). A preset value of 0.6 is recommended for most cases. <br>
-            2.5. To use the motion brush for restraining the control area of the trajectory, click to add masks on the "Add Motion Brush Here" image. The motion brush restricts the optical flow area derived from the trajectory whose starting point is within the motion brush. The displayed optical flow image will change correspondingly. Adjust the motion brush radius using the "Motion Brush Radius" bar. <br>
-        3. Click the "Run" button to animate the image according to the path. <br>
-        """
-    )
-
-    target_size = 512
-    DragNUWA_net = Drag(target_size, target_size)
-    first_frame_path = gr.State()
-    tracking_points = gr.State([])
-    motion_brush_points = gr.State([])
-    motion_brush_mask = gr.State()
-    motion_brush_viz = gr.State()
-    inference_batch_size = gr.State(1)
-
+    @spaces.GPU(duration=100)
     def preprocess_image(image):
+        
+        if pipeline is None or cmp is None:
+            pipeline, cmp = init_models()
 
         image_pil = image2pil(image.name)
         raw_w, raw_h = image_pil.size
 
         max_edge = min(raw_w, raw_h)
-        resize_ratio = target_size / max_edge
+        resize_ratio = width / max_edge
 
         image_pil = image_pil.resize((round(raw_w * resize_ratio), round(raw_h * resize_ratio)), Image.BILINEAR)
 
@@ -676,8 +664,8 @@ with gr.Blocks() as demo:
 
         image_pil = transforms.CenterCrop((crop_h, crop_w))(image_pil.convert('RGB'))
 
-        DragNUWA_net.width = crop_w
-        DragNUWA_net.height = crop_h
+        width = crop_w
+        height = crop_h
 
         id = str(time.time()).split('.')[0]
         os.makedirs(os.path.join(output_dir_video, str(id)), exist_ok=True)
@@ -722,7 +710,7 @@ with gr.Blocks() as demo:
         transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
         trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
 
-        viz_flow = DragNUWA_net.get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
+        viz_flow = get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
 
         return tracking_points, trajectory_map, viz_flow
     
@@ -742,7 +730,7 @@ with gr.Blocks() as demo:
         transparent_layer_pil = Image.fromarray(transparent_layer.astype(np.uint8))
         motion_map = Image.alpha_composite(transparent_background, transparent_layer_pil)
 
-        viz_flow = DragNUWA_net.get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
+        viz_flow = get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
 
         return motion_brush_mask, transparent_layer, motion_map, viz_flow
 
@@ -778,7 +766,7 @@ with gr.Blocks() as demo:
         transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
         trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
 
-        viz_flow = DragNUWA_net.get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
+        viz_flow = get_cmp_flow_from_tracking_points(tracking_points, motion_brush_mask, first_frame_path)
 
         return tracking_points, trajectory_map, viz_flow
 
@@ -833,6 +821,6 @@ with gr.Blocks() as demo:
 
     input_image_mask.select(add_motion_brushes, [motion_brush_points, motion_brush_mask, motion_brush_viz, first_frame_path, brush_radius, tracking_points], [motion_brush_mask, motion_brush_viz, input_image_mask, viz_flow])
 
-    run_button.click(DragNUWA_net.run, [first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale], [hint_image, output_video, output_flow, output_video_mp4, output_flow_mp4])
+    run_button.click(run, [first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale], [hint_image, output_video, output_flow, output_video_mp4, output_flow_mp4])
 
     demo.launch()