Update

app.py

@@ -10,9 +10,7 @@ import numpy as np
 import spaces
 import torch
 import torchvision
-from diffusers.utils.import_utils import is_xformers_available
 from huggingface_hub import snapshot_download
-from packaging import version
 from PIL import Image
 from scipy.interpolate import PchipInterpolator
 
@@ -39,55 +37,40 @@ snapshot_download(
 )
 
 
-def get_args():
-    import argparse
+model_id = "checkpoints/framer_512x320"
+device = "cuda"
+dtype = torch.float16
 
-    parser = argparse.ArgumentParser()
+OUTPUT_DIR = "gradio_demo/outputs"
+HEIGHT = 320
+WIDTH = 512
+MODEL_LENGTH = 14
+USE_SIFT = False
 
-    parser.add_argument("--min_guidance_scale", type=float, default=1.0)
-    parser.add_argument("--max_guidance_scale", type=float, default=3.0)
-    parser.add_argument("--middle_max_guidance", type=int, default=0, choices=[0, 1])
-    parser.add_argument("--with_control", type=int, default=1, choices=[0, 1])
 
-    parser.add_argument("--controlnet_cond_scale", type=float, default=1.0)
-
-    parser.add_argument(
-        "--dataset",
-        type=str,
-        default="videoswap",
-    )
-
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="checkpoints/framer_512x320",
-        help="Path to model.",
-    )
-
-    parser.add_argument("--output_dir", type=str, default="gradio_demo/outputs", help="Path to the output video.")
-
-    parser.add_argument("--seed", type=int, default=42, help="random seed.")
-
-    parser.add_argument("--noise_aug", type=float, default=0.02)
-
-    parser.add_argument("--num_frames", type=int, default=14)
-    parser.add_argument("--frame_interval", type=int, default=2)
-
-    parser.add_argument("--width", type=int, default=512)
-    parser.add_argument("--height", type=int, default=320)
-
-    parser.add_argument(
-        "--num_workers",
-        type=int,
-        default=0,
-        help=(
-            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
-        ),
-    )
-
-    args = parser.parse_args()
+unet = UNetSpatioTemporalConditionModel.from_pretrained(
+    os.path.join(model_id, "unet"),
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    custom_resume=True,
+)
+unet = unet.to(device, dtype)
 
-    return args
+controlnet = ControlNetSVDModel.from_pretrained(
+    os.path.join(model_id, "controlnet"),
+)
+controlnet = controlnet.to(device, dtype)
+
+pipe = StableVideoDiffusionInterpControlPipeline.from_pretrained(
+    "checkpoints/stable-video-diffusion-img2vid-xt",
+    unet=unet,
+    controlnet=controlnet,
+    low_cpu_mem_usage=False,
+    torch_dtype=torch.float16,
+    variant="fp16",
+    local_files_only=True,
+)
+pipe.to(device)
 
 
 def interpolate_trajectory(points, n_points):
@@ -164,7 +147,7 @@ def get_vis_image(
         vis_img = new_img.copy()
         # ids_embedding = torch.zeros((target_size[0], target_size[1], 320))
 
-        if idxx >= args.num_frames:
+        if idxx >= num_frames:
             break
 
         # for cc, (mask, trajectory, radius) in enumerate(zip(mask_list, trajectory_list, radius_list)):
@@ -363,187 +346,6 @@ def validate_and_convert_image(image, target_size=(512, 512)):
     return image
 
 
-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
-
-        unet = UNetSpatioTemporalConditionModel.from_pretrained(
-            os.path.join(args.model, "unet"),
-            torch_dtype=torch.float16,
-            low_cpu_mem_usage=True,
-            custom_resume=True,
-        )
-        unet = unet.to(device, dtype)
-
-        controlnet = ControlNetSVDModel.from_pretrained(
-            os.path.join(args.model, "controlnet"),
-        )
-        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")
-
-        pipe = StableVideoDiffusionInterpControlPipeline.from_pretrained(
-            "checkpoints/stable-video-diffusion-img2vid-xt",
-            unet=unet,
-            controlnet=controlnet,
-            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()
-
-        self.height = height
-        self.width = width
-        self.args = args
-        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
-
-        # load_image
-        image = Image.open(first_frame_path).convert("RGB")
-        width, height = image.size
-        image = image.resize((self.width, self.height))
-
-        image_end = Image.open(last_frame_path).convert("RGB")
-        image_end = image_end.resize((self.width, self.height))
-
-        input_all_points = tracking_points
-
-        sift_track_update = False
-        anchor_points_flag = None
-
-        if (len(input_all_points) == 0) and self.use_sift:
-            sift_track_update = True
-            controlnet_cond_scale = 0.5
-
-            from models_diffusers.sift_match import interpolate_trajectory as sift_interpolate_trajectory
-            from models_diffusers.sift_match import sift_match
-
-            output_file_sift = os.path.join(args.output_dir, "sift.png")
-
-            # (f, topk, 2), f=2 (before interpolation)
-            pred_tracks = sift_match(
-                image,
-                image_end,
-                thr=0.5,
-                topk=5,
-                method="random",
-                output_path=output_file_sift,
-            )
-
-            if pred_tracks is not None:
-                # interpolate the tracks, following draganything gradio demo
-                pred_tracks = sift_interpolate_trajectory(pred_tracks, num_frames=self.model_length)
-
-                anchor_points_flag = torch.zeros((self.model_length, pred_tracks.shape[1])).to(pred_tracks.device)
-                anchor_points_flag[0] = 1
-                anchor_points_flag[-1] = 1
-
-                pred_tracks = pred_tracks.permute(1, 0, 2)  # (num_points, num_frames, 2)
-
-        else:
-
-            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
-            ]
-
-            # a list of num_tracks tuples, each tuple contains a track with several points, represented as (x, y)
-            # in image w & h scale
-
-            for idx, splited_track in enumerate(resized_all_points):
-                if len(splited_track) == 0:
-                    warnings.warn("running without point trajectory control")
-                    continue
-
-                if len(splited_track) == 1:  # stationary point
-                    displacement_point = tuple([splited_track[0][0] + 1, splited_track[0][1] + 1])
-                    splited_track = tuple([splited_track[0], displacement_point])
-                # interpolate the track
-                splited_track = interpolate_trajectory(splited_track, self.model_length)
-                splited_track = splited_track[: self.model_length]
-                resized_all_points[idx] = splited_track
-
-            pred_tracks = torch.tensor(resized_all_points)  # (num_points, num_frames, 2)
-
-        vis_images = get_vis_image(
-            target_size=(self.args.height, self.args.width),
-            points=pred_tracks,
-            num_frames=self.model_length,
-        )
-
-        if len(pred_tracks.shape) != 3:
-            print("pred_tracks.shape", pred_tracks.shape)
-            with_control = False
-            controlnet_cond_scale = 0.0
-        else:
-            with_control = True
-            pred_tracks = pred_tracks.permute(1, 0, 2).to(self.device, self.dtype)  # (num_frames, num_points, 2)
-
-        point_embedding = None
-        video_frames = self.pipeline(
-            image,
-            image_end,
-            # trajectory control
-            with_control=with_control,
-            point_tracks=pred_tracks,
-            point_embedding=point_embedding,
-            with_id_feature=False,
-            controlnet_cond_scale=controlnet_cond_scale,
-            # others
-            num_frames=14,
-            width=width,
-            height=height,
-            # decode_chunk_size=8,
-            # generator=generator,
-            motion_bucket_id=motion_bucket_id,
-            fps=7,
-            num_inference_steps=30,
-            # track
-            sift_track_update=sift_track_update,
-            anchor_points_flag=anchor_points_flag,
-        ).frames[0]
-
-        vis_images = [cv2.applyColorMap(np.array(img).astype(np.uint8), cv2.COLORMAP_JET) for img in vis_images]
-        vis_images = [cv2.cvtColor(np.array(img).astype(np.uint8), cv2.COLOR_BGR2RGB) for img in vis_images]
-        vis_images = [Image.fromarray(img) for img in vis_images]
-
-        # video_frames = [img for sublist in video_frames for img in sublist]
-        val_save_dir = os.path.join(args.output_dir, "vis_gif.gif")
-        save_gifs_side_by_side(
-            video_frames,
-            vis_images[: self.model_length],
-            val_save_dir,
-            target_size=(self.width, self.height),
-            duration=110,
-            point_tracks=pred_tracks,
-        )
-
-        return val_save_dir
-
-
 def reset_states(first_frame_path, last_frame_path, tracking_points):
     first_frame_path = None
     last_frame_path = None
@@ -561,7 +363,7 @@ def preprocess_image(image):
     # image_pil = transforms.CenterCrop((320, 512))(image_pil.convert('RGB'))
     image_pil = image_pil.resize((512, 320), Image.BILINEAR)
 
-    first_frame_path = os.path.join(args.output_dir, f"first_frame_{str(uuid.uuid4())[:4]}.png")
+    first_frame_path = os.path.join(OUTPUT_DIR, f"first_frame_{str(uuid.uuid4())[:4]}.png")
 
     image_pil.save(first_frame_path)
 
@@ -578,7 +380,7 @@ def preprocess_image_end(image_end):
     # image_end_pil = transforms.CenterCrop((320, 512))(image_end_pil.convert('RGB'))
     image_end_pil = image_end_pil.resize((512, 320), Image.BILINEAR)
 
-    last_frame_path = os.path.join(args.output_dir, f"last_frame_{str(uuid.uuid4())[:4]}.png")
+    last_frame_path = os.path.join(OUTPUT_DIR, f"last_frame_{str(uuid.uuid4())[:4]}.png")
 
     image_end_pil.save(last_frame_path)
 
@@ -692,7 +494,7 @@ def add_tracking_points(
     transparent_layer = 0
     for idx, track in enumerate(tracking_points):
         # mask = cv2.imread(
-        #     os.path.join(args.output_dir, f"mask_{idx+1}.jpg")
+        #     os.path.join(OUTPUT_DIR, f"mask_{idx+1}.jpg")
         # )
         mask = np.zeros((320, 512, 3))
         color = color_list[idx + 1]
@@ -737,10 +539,136 @@ def add_tracking_points(
     return tracking_points, trajectory_map, trajectory_map_end
 
 
+@spaces.GPU
+def run(first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id):
+    original_width, original_height = 512, 320  # TODO
+
+    # load_image
+    image = Image.open(first_frame_path).convert("RGB")
+    width, height = image.size
+    image = image.resize((WIDTH, HEIGHT))
+
+    image_end = Image.open(last_frame_path).convert("RGB")
+    image_end = image_end.resize((WIDTH, HEIGHT))
+
+    input_all_points = tracking_points
+
+    sift_track_update = False
+    anchor_points_flag = None
+
+    if (len(input_all_points) == 0) and USE_SIFT:
+        sift_track_update = True
+        controlnet_cond_scale = 0.5
+
+        from models_diffusers.sift_match import interpolate_trajectory as sift_interpolate_trajectory
+        from models_diffusers.sift_match import sift_match
+
+        output_file_sift = os.path.join(OUTPUT_DIR, "sift.png")
+
+        # (f, topk, 2), f=2 (before interpolation)
+        pred_tracks = sift_match(
+            image,
+            image_end,
+            thr=0.5,
+            topk=5,
+            method="random",
+            output_path=output_file_sift,
+        )
+
+        if pred_tracks is not None:
+            # interpolate the tracks, following draganything gradio demo
+            pred_tracks = sift_interpolate_trajectory(pred_tracks, num_frames=MODEL_LENGTH)
+
+            anchor_points_flag = torch.zeros((MODEL_LENGTH, pred_tracks.shape[1])).to(pred_tracks.device)
+            anchor_points_flag[0] = 1
+            anchor_points_flag[-1] = 1
+
+            pred_tracks = pred_tracks.permute(1, 0, 2)  # (num_points, num_frames, 2)
+
+    else:
+
+        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
+        ]
+
+        # a list of num_tracks tuples, each tuple contains a track with several points, represented as (x, y)
+        # in image w & h scale
+
+        for idx, splited_track in enumerate(resized_all_points):
+            if len(splited_track) == 0:
+                warnings.warn("running without point trajectory control")
+                continue
+
+            if len(splited_track) == 1:  # stationary point
+                displacement_point = tuple([splited_track[0][0] + 1, splited_track[0][1] + 1])
+                splited_track = tuple([splited_track[0], displacement_point])
+            # interpolate the track
+            splited_track = interpolate_trajectory(splited_track, MODEL_LENGTH)
+            splited_track = splited_track[:MODEL_LENGTH]
+            resized_all_points[idx] = splited_track
+
+        pred_tracks = torch.tensor(resized_all_points)  # (num_points, num_frames, 2)
+
+    vis_images = get_vis_image(
+        target_size=(HEIGHT, WIDTH),
+        points=pred_tracks,
+        num_frames=MODEL_LENGTH,
+    )
+
+    if len(pred_tracks.shape) != 3:
+        print("pred_tracks.shape", pred_tracks.shape)
+        with_control = False
+        controlnet_cond_scale = 0.0
+    else:
+        with_control = True
+        pred_tracks = pred_tracks.permute(1, 0, 2).to(device, dtype)  # (num_frames, num_points, 2)
+
+    point_embedding = None
+    video_frames = pipe(
+        image,
+        image_end,
+        # trajectory control
+        with_control=with_control,
+        point_tracks=pred_tracks,
+        point_embedding=point_embedding,
+        with_id_feature=False,
+        controlnet_cond_scale=controlnet_cond_scale,
+        # others
+        num_frames=14,
+        width=width,
+        height=height,
+        # decode_chunk_size=8,
+        # generator=generator,
+        motion_bucket_id=motion_bucket_id,
+        fps=7,
+        num_inference_steps=30,
+        # track
+        sift_track_update=sift_track_update,
+        anchor_points_flag=anchor_points_flag,
+    ).frames[0]
+
+    vis_images = [cv2.applyColorMap(np.array(img).astype(np.uint8), cv2.COLORMAP_JET) for img in vis_images]
+    vis_images = [cv2.cvtColor(np.array(img).astype(np.uint8), cv2.COLOR_BGR2RGB) for img in vis_images]
+    vis_images = [Image.fromarray(img) for img in vis_images]
+
+    # video_frames = [img for sublist in video_frames for img in sublist]
+    val_save_dir = os.path.join(OUTPUT_DIR, "vis_gif.gif")
+    save_gifs_side_by_side(
+        video_frames,
+        vis_images[:MODEL_LENGTH],
+        val_save_dir,
+        target_size=(WIDTH, HEIGHT),
+        duration=110,
+        point_tracks=pred_tracks,
+    )
+
+    return val_save_dir
+
+
 if __name__ == "__main__":
 
-    args = get_args()
-    ensure_dirname(args.output_dir)
+    ensure_dirname(OUTPUT_DIR)
 
     color_list = []
     for i in range(20):
@@ -771,8 +699,6 @@ if __name__ == "__main__":
                     3. Interpolate the images (according the path) with a click on "Run" button. <br>"""
         )
 
-        # device, args, height, width, model_length
-        Framer = Drag("cuda", args, 320, 512, 14)
         first_frame_path = gr.State()
         last_frame_path = gr.State()
         tracking_points = gr.State([])
@@ -898,7 +824,7 @@ if __name__ == "__main__":
         )
 
         run_button.click(
-            fn=Framer.run,
+            fn=run,
             inputs=[first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id],
             outputs=output_video,
         )