ZeroGPU

README.md

@@ -4,7 +4,7 @@ emoji: 🏃
 colorFrom: gray
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.41.0
+sdk_version: 5.5.0
 python_version: 3.8.9
 app_file: app.py
 pinned: false

app.py

@@ -1,94 +1,76 @@
-import spaces
 import datetime
+import os
+import sys
 import uuid
-from PIL import Image
-import numpy as np
-import cv2
-from scipy.interpolate import interp1d, PchipInterpolator
-from packaging import version
+import warnings
 
+import cv2
+import gradio as gr
+import numpy as np
+import spaces
 import torch
 import torchvision
-import gradio as gr
-# from moviepy.editor import *
-from diffusers.utils.import_utils import is_xformers_available
-from diffusers.utils import load_image, export_to_video, export_to_gif
+from huggingface_hub import snapshot_download
+from PIL import Image
+from scipy.interpolate import PchipInterpolator
 
-import os
-import sys
 sys.path.insert(0, os.getcwd())
+
+from gradio_demo.utils_drag import *
 from models_diffusers.controlnet_svd import ControlNetSVDModel
 from models_diffusers.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
 from pipelines.pipeline_stable_video_diffusion_interp_control import StableVideoDiffusionInterpControlPipeline
-from gradio_demo.utils_drag import *
 
-import warnings
 print("gr file", gr.__file__)
 
-from huggingface_hub import hf_hub_download, snapshot_download
 
 os.makedirs("checkpoints", exist_ok=True)
 
 snapshot_download(
     "wwen1997/framer_512x320",
     local_dir="checkpoints/framer_512x320",
-    token=os.environ["TOKEN"],
 )
 
 snapshot_download(
     "stabilityai/stable-video-diffusion-img2vid-xt",
     local_dir="checkpoints/stable-video-diffusion-img2vid-xt",
-    token=os.environ["TOKEN"],
 )
 
 
-def get_args():
-    import argparse
-    parser = argparse.ArgumentParser()
-
-    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.")
+model_id = "checkpoints/framer_512x320"
+device = "cuda"
+dtype = torch.float16
 
-    parser.add_argument("--noise_aug", type=float, default=0.02)
+OUTPUT_DIR = "gradio_demo/outputs"
+HEIGHT = 320
+WIDTH = 512
+MODEL_LENGTH = 14
+USE_SIFT = False
 
-    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):
@@ -113,27 +95,34 @@ def interpolate_trajectory(points, n_points):
 
 def gen_gaussian_heatmap(imgSize=200):
     circle_img = np.zeros((imgSize, imgSize), np.float32)
-    circle_mask = cv2.circle(circle_img, (imgSize//2, imgSize//2), imgSize//2, 1, -1)
+    circle_mask = cv2.circle(circle_img, (imgSize // 2, imgSize // 2), imgSize // 2, 1, -1)
 
     isotropicGrayscaleImage = np.zeros((imgSize, imgSize), np.float32)
 
     for i in range(imgSize):
         for j in range(imgSize):
-            isotropicGrayscaleImage[i, j] = 1 / 2 / np.pi / (40 ** 2) * np.exp(
-                -1 / 2 * ((i - imgSize / 2) ** 2 / (40 ** 2) + (j - imgSize / 2) ** 2 / (40 ** 2)))
+            isotropicGrayscaleImage[i, j] = (
+                1
+                / 2
+                / np.pi
+                / (40**2)
+                * np.exp(-1 / 2 * ((i - imgSize / 2) ** 2 / (40**2) + (j - imgSize / 2) ** 2 / (40**2)))
+            )
 
     isotropicGrayscaleImage = isotropicGrayscaleImage * circle_mask
     isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage)).astype(np.float32)
-    isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage)*255).astype(np.uint8)
+    isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage) * 255).astype(np.uint8)
 
     return isotropicGrayscaleImage
 
 
 def get_vis_image(
-        target_size=(512 , 512), points=None,  side=20,
-        num_frames=14,
-        # original_size=(512 , 512), args="", first_frame=None, is_mask = False, model_id=None,
-    ):
+    target_size=(512, 512),
+    points=None,
+    side=20,
+    num_frames=14,
+    # original_size=(512 , 512), args="", first_frame=None, is_mask = False, model_id=None,
+):
 
     # images = []
     vis_images = []
@@ -141,13 +130,13 @@ def get_vis_image(
 
     trajectory_list = []
     radius_list = []
-    
+
     for index, point in enumerate(points):
         trajectories = [[int(i[0]), int(i[1])] for i in point]
         trajectory_list.append(trajectories)
 
         radius = 20
-        radius_list.append(radius)  
+        radius_list.append(radius)
 
     if len(trajectory_list) == 0:
         vis_images = [Image.fromarray(np.zeros(target_size, np.uint8)) for _ in range(num_frames)]
@@ -157,33 +146,39 @@ def get_vis_image(
         new_img = np.zeros(target_size, np.uint8)
         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)):
         for cc, (trajectory, radius) in enumerate(zip(trajectory_list, radius_list)):
-            
+
             center_coordinate = trajectory[idxx]
             trajectory_ = trajectory[:idxx]
             side = min(radius, 50)
- 
-            y1 = max(center_coordinate[1] - side,0)
+
+            y1 = max(center_coordinate[1] - side, 0)
             y2 = min(center_coordinate[1] + side, target_size[0] - 1)
             x1 = max(center_coordinate[0] - side, 0)
             x2 = min(center_coordinate[0] + side, target_size[1] - 1)
-            
-            if x2-x1>3 and y2-y1>3:
-                need_map = cv2.resize(heatmap, (x2-x1, y2-y1))
+
+            if x2 - x1 > 3 and y2 - y1 > 3:
+                need_map = cv2.resize(heatmap, (x2 - x1, y2 - y1))
                 new_img[y1:y2, x1:x2] = need_map.copy()
-                
+
                 if cc >= 0:
-                    vis_img[y1:y2,x1:x2] = need_map.copy()
+                    vis_img[y1:y2, x1:x2] = need_map.copy()
                     if len(trajectory_) == 1:
                         vis_img[trajectory_[0][1], trajectory_[0][0]] = 255
                     else:
-                        for itt in range(len(trajectory_)-1):
-                            cv2.line(vis_img, (trajectory_[itt][0], trajectory_[itt][1]), (trajectory_[itt+1][0], trajectory_[itt+1][1]), (255, 255, 255), 3)
+                        for itt in range(len(trajectory_) - 1):
+                            cv2.line(
+                                vis_img,
+                                (trajectory_[itt][0], trajectory_[itt][1]),
+                                (trajectory_[itt + 1][0], trajectory_[itt + 1][1]),
+                                (255, 255, 255),
+                                3,
+                            )
 
         img = new_img
 
@@ -194,7 +189,7 @@ def get_vis_image(
         elif len(img.shape) == 3 and img.shape[2] == 3:  # Color image in BGR format
             img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
             vis_img = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
-            
+
         # Convert the numpy array to a PIL image
         # pil_img = Image.fromarray(img)
         # images.append(pil_img)
@@ -215,7 +210,7 @@ def frames_to_video(frames_folder, output_video_path, fps=7):
         video.append(frame)
 
     video = torch.stack(video)
-    video = rearrange(video, 'T C H W -> T H W C')
+    video = rearrange(video, "T C H W -> T H W C")
     torchvision.io.write_video(output_video_path, video, fps=fps)
 
 
@@ -223,11 +218,12 @@ def save_gifs_side_by_side(
     batch_output,
     validation_control_images,
     output_folder,
-    target_size=(512 , 512),
+    target_size=(512, 512),
     duration=200,
     point_tracks=None,
 ):
     flattened_batch_output = batch_output
+
     def create_gif(image_list, gif_path, duration=100):
         pil_images = [validate_and_convert_image(img, target_size=target_size) for img in image_list]
         pil_images = [img for img in pil_images if img is not None]
@@ -243,7 +239,7 @@ def save_gifs_side_by_side(
             tmp_frame_path = os.path.join(tmp_folder, f"{idx}.png")
             pil_image.save(tmp_frame_path)
             tmp_frame_list.append(tmp_frame_path)
-        
+
         # also save as mp4
         output_video_path = gif_path.replace(".gif", ".mp4")
         frames_to_video(tmp_folder, output_video_path, fps=7)
@@ -286,25 +282,25 @@ def save_gifs_side_by_side(
         if output_path.endswith(".mp4"):
             video = [torchvision.transforms.functional.pil_to_tensor(frame) for frame in frames]
             video = torch.stack(video)
-            video = rearrange(video, 'T C H W -> T H W C')
+            video = rearrange(video, "T C H W -> T H W C")
             torchvision.io.write_video(output_path, video, fps=7)
             print(f"Saved video to {output_path}")
         else:
             frames[0].save(output_path, save_all=True, append_images=frames[1:], loop=0, duration=duration)
-        
+
     # Helper function to concatenate images horizontally
     def get_concat_h(im1, im2, gap=10):
         # # img first, heatmap second
         # im1, im2 = im2, im1
 
-        dst = Image.new('RGB', (im1.width + im2.width + gap, max(im1.height, im2.height)), (255, 255, 255))
+        dst = Image.new("RGB", (im1.width + im2.width + gap, max(im1.height, im2.height)), (255, 255, 255))
         dst.paste(im1, (0, 0))
         dst.paste(im2, (im1.width + gap, 0))
         return dst
 
     # Helper function to concatenate images vertically
     def get_concat_v(im1, im2):
-        dst = Image.new('RGB', (max(im1.width, im2.width), im1.height + im2.height))
+        dst = Image.new("RGB", (max(im1.width, im2.width), im1.height + im2.height))
         dst.paste(im1, (0, 0))
         dst.paste(im2, (0, im1.height))
         return dst
@@ -325,7 +321,7 @@ def save_gifs_side_by_side(
 
 
 # Define functions
-def validate_and_convert_image(image, target_size=(512 , 512)):
+def validate_and_convert_image(image, target_size=(512, 512)):
     if image is None:
         print("Encountered a None image")
         return None
@@ -346,192 +342,12 @@ def validate_and_convert_image(image, target_size=(512 , 512)):
     else:
         print("Image is not a PIL Image or a PyTorch tensor")
         return None
-    
-    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.constructor_args['value']
-
-        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 sift_match
-            from models_diffusers.sift_match import interpolate_trajectory as sift_interpolate_trajectory
-
-            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
 
+    return image
 
-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([])
 
-    return first_frame_path, last_frame_path, tracking_points
+def reset_states():
+    return None, None, None, None, None, []
 
 
 def preprocess_image(image):
@@ -544,11 +360,11 @@ 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)
 
-    return first_frame_path, first_frame_path, gr.State([])
+    return first_frame_path, first_frame_path, []
 
 
 def preprocess_image_end(image_end):
@@ -561,37 +377,52 @@ 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)
 
-    return last_frame_path, last_frame_path, gr.State([])
+    return last_frame_path, last_frame_path, []
 
 
 def add_drag(tracking_points):
-    tracking_points.constructor_args['value'].append([])
+    if not tracking_points or tracking_points[-1]:
+        tracking_points.append([])
     return tracking_points
 
 
 def delete_last_drag(tracking_points, first_frame_path, last_frame_path):
-    tracking_points.constructor_args['value'].pop()
-    transparent_background = Image.open(first_frame_path).convert('RGBA')
-    transparent_background_end = Image.open(last_frame_path).convert('RGBA')
+    if tracking_points:
+        tracking_points.pop()
+    transparent_background = Image.open(first_frame_path).convert("RGBA")
+    transparent_background_end = Image.open(last_frame_path).convert("RGBA")
     w, h = transparent_background.size
     transparent_layer = np.zeros((h, w, 4))
 
-    for track in tracking_points.constructor_args['value']:
+    for track in tracking_points:
         if len(track) > 1:
-            for i in range(len(track)-1):
+            for i in range(len(track) - 1):
                 start_point = track[i]
-                end_point = track[i+1]
+                end_point = track[i + 1]
                 vx = end_point[0] - start_point[0]
                 vy = end_point[1] - start_point[1]
                 arrow_length = np.sqrt(vx**2 + vy**2)
-                if i == len(track)-2:
-                    cv2.arrowedLine(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                if i == len(track) - 2:
+                    cv2.arrowedLine(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        (255, 0, 0, 255),
+                        2,
+                        tipLength=8 / arrow_length,
+                    )
                 else:
-                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+                    cv2.line(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        (255, 0, 0, 255),
+                        2,
+                    )
         else:
             cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
 
@@ -603,24 +434,40 @@ def delete_last_drag(tracking_points, first_frame_path, last_frame_path):
 
 
 def delete_last_step(tracking_points, first_frame_path, last_frame_path):
-    tracking_points.constructor_args['value'][-1].pop()
-    transparent_background = Image.open(first_frame_path).convert('RGBA')
-    transparent_background_end = Image.open(last_frame_path).convert('RGBA')
+    if tracking_points and tracking_points[-1]:
+        tracking_points[-1].pop()
+    transparent_background = Image.open(first_frame_path).convert("RGBA")
+    transparent_background_end = Image.open(last_frame_path).convert("RGBA")
     w, h = transparent_background.size
     transparent_layer = np.zeros((h, w, 4))
 
-    for track in tracking_points.constructor_args['value']:
+    for track in tracking_points:
+        if not track:
+            continue
         if len(track) > 1:
-            for i in range(len(track)-1):
+            for i in range(len(track) - 1):
                 start_point = track[i]
-                end_point = track[i+1]
+                end_point = track[i + 1]
                 vx = end_point[0] - start_point[0]
                 vy = end_point[1] - start_point[1]
                 arrow_length = np.sqrt(vx**2 + vy**2)
-                if i == len(track)-2:
-                    cv2.arrowedLine(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                if i == len(track) - 2:
+                    cv2.arrowedLine(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        (255, 0, 0, 255),
+                        2,
+                        tipLength=8 / arrow_length,
+                    )
                 else:
-                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+                    cv2.line(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        (255, 0, 0, 255),
+                        2,
+                    )
         else:
             cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
 
@@ -631,34 +478,51 @@ def delete_last_step(tracking_points, first_frame_path, last_frame_path):
     return tracking_points, trajectory_map, trajectory_map_end
 
 
-def add_tracking_points(tracking_points, first_frame_path, last_frame_path, evt: gr.SelectData):  # SelectData is a subclass of EventData
+def add_tracking_points(
+    tracking_points, first_frame_path, last_frame_path, evt: gr.SelectData
+):  # SelectData is a subclass of EventData
     print(f"You selected {evt.value} at {evt.index} from {evt.target}")
-    tracking_points.constructor_args['value'][-1].append(evt.index)
+    if not tracking_points:
+        tracking_points = [[]]
+    tracking_points[-1].append(evt.index)
 
-    transparent_background = Image.open(first_frame_path).convert('RGBA')
-    transparent_background_end = Image.open(last_frame_path).convert('RGBA')
+    transparent_background = Image.open(first_frame_path).convert("RGBA")
+    transparent_background_end = Image.open(last_frame_path).convert("RGBA")
 
     w, h = transparent_background.size
     transparent_layer = 0
-    for idx, track in enumerate(tracking_points.constructor_args['value']):
+    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]
+        color = color_list[idx + 1]
         transparent_layer = mask[:, :, 0].reshape(h, w, 1) * color.reshape(1, 1, -1) + transparent_layer
 
         if len(track) > 1:
-            for i in range(len(track)-1):
+            for i in range(len(track) - 1):
                 start_point = track[i]
-                end_point = track[i+1]
+                end_point = track[i + 1]
                 vx = end_point[0] - start_point[0]
                 vy = end_point[1] - start_point[1]
                 arrow_length = np.sqrt(vx**2 + vy**2)
-                if i == len(track)-2:
-                    cv2.arrowedLine(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2, tipLength=8 / arrow_length)
+                if i == len(track) - 2:
+                    cv2.arrowedLine(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        (255, 0, 0, 255),
+                        2,
+                        tipLength=8 / arrow_length,
+                    )
                 else:
-                    cv2.line(transparent_layer, tuple(start_point), tuple(end_point), (255, 0, 0, 255), 2,)
+                    cv2.line(
+                        transparent_layer,
+                        tuple(start_point),
+                        tuple(end_point),
+                        (255, 0, 0, 255),
+                        2,
+                    )
         else:
             cv2.circle(transparent_layer, tuple(track[0]), 5, (255, 0, 0, 255), -1)
 
@@ -674,26 +538,162 @@ def add_tracking_points(tracking_points, first_frame_path, last_frame_path, evt:
     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,
+    progress=gr.Progress(track_tqdm=True),
+):
+    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):
-        color = np.concatenate([np.random.random(4)*255], axis=0)
+        color = np.concatenate([np.random.random(4) * 255], axis=0)
         color_list.append(color)
 
     with gr.Blocks() as demo:
         gr.Markdown("""<h1 align="center">Framer: Interactive Frame Interpolation</h1><br>""")
-    
-        gr.Markdown("""Gradio Demo for <a href='https://arxiv.org/abs/2410.18978'><b>Framer: Interactive Frame Interpolation</b></a>.<br>
+
+        gr.Markdown(
+            """Gradio Demo for <a href='https://arxiv.org/abs/2410.18978'><b>Framer: Interactive Frame Interpolation</b></a>.<br>
                     Github Repo can be found at https://github.com/aim-uofa/Framer<br>
-                    The template is inspired by DragAnything.""")
-    
+                    The template is inspired by DragAnything."""
+        )
+
         gr.Image(label="Framer: Interactive Frame Interpolation", value="assets/demos.gif", height=432, width=768)
-    
-        gr.Markdown("""## Usage: <br>
+
+        gr.Markdown(
+            """## Usage: <br>
                     1. Upload images<br>
                     &ensp;  1.1  Upload the start image via the "Upload Start Image" button.<br>
                     &ensp;  1.2. Upload the end image via the "Upload End Image" button.<br>
@@ -702,14 +702,13 @@ if __name__ == "__main__":
                     &ensp;  2.2. You can click several points on either start or end image to forms a path.<br>
                     &ensp;  2.3. Click "Delete last drag" to delete the whole lastest path.<br>
                     &ensp;  2.4. Click "Delete last step" to delete the lastest clicked control point.<br>
-                    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)
+                    3. Interpolate the images (according the path) with a click on "Run" button. <br>"""
+        )
+
         first_frame_path = gr.State()
         last_frame_path = gr.State()
         tracking_points = gr.State([])
-    
+
         with gr.Row():
             with gr.Column(scale=1):
                 image_upload_button = gr.UploadButton(label="Upload Start Image", file_types=["image"])
@@ -720,7 +719,7 @@ if __name__ == "__main__":
                 run_button = gr.Button(value="Run")
                 delete_last_drag_button = gr.Button(value="Delete last drag")
                 delete_last_step_button = gr.Button(value="Delete last step")
-    
+
             with gr.Column(scale=7):
                 with gr.Row():
                     with gr.Column(scale=6):
@@ -731,7 +730,7 @@ if __name__ == "__main__":
                             width=512,
                             sources=[],
                         )
-    
+
                     with gr.Column(scale=6):
                         input_image_end = gr.Image(
                             label="end frame",
@@ -740,36 +739,36 @@ if __name__ == "__main__":
                             width=512,
                             sources=[],
                         )
-    
+
         with gr.Row():
             with gr.Column(scale=1):
-    
+
                 controlnet_cond_scale = gr.Slider(
-                    label='Control Scale', 
-                    minimum=0.0, 
-                    maximum=10, 
-                    step=0.1, 
+                    label="Control Scale",
+                    minimum=0.0,
+                    maximum=10,
+                    step=0.1,
                     value=1.0,
                 )
-    
+
                 motion_bucket_id = gr.Slider(
-                    label='Motion Bucket', 
-                    minimum=1, 
-                    maximum=180, 
-                    step=1, 
+                    label="Motion Bucket",
+                    minimum=1,
+                    maximum=180,
+                    step=1,
                     value=100,
                 )
-    
+
             with gr.Column(scale=5):
                 output_video = gr.Image(
                     label="Output Video",
                     height=320,
                     width=1152,
                 )
-    
-    
+
         with gr.Row():
-            gr.Markdown("""
+            gr.Markdown(
+                """
                 ## Citation
                 ```bibtex
                 @article{wang2024framer,
@@ -779,24 +778,55 @@ if __name__ == "__main__":
                   year={2024}
                 }
                 ```
-                """)
-    
-        image_upload_button.upload(preprocess_image, image_upload_button, [input_image, first_frame_path, tracking_points])
-    
-        image_end_upload_button.upload(preprocess_image_end, image_end_upload_button, [input_image_end, last_frame_path, tracking_points])
-    
-        add_drag_button.click(add_drag, tracking_points, [tracking_points, ])
-    
-        delete_last_drag_button.click(delete_last_drag, [tracking_points, first_frame_path, last_frame_path], [tracking_points, input_image, input_image_end])
-    
-        delete_last_step_button.click(delete_last_step, [tracking_points, first_frame_path, last_frame_path], [tracking_points, input_image, input_image_end])
-    
-        reset_button.click(reset_states, [first_frame_path, last_frame_path, tracking_points], [first_frame_path, last_frame_path, tracking_points])
-    
-        input_image.select(add_tracking_points, [tracking_points, first_frame_path, last_frame_path], [tracking_points, input_image, input_image_end])
-    
-        input_image_end.select(add_tracking_points, [tracking_points, first_frame_path, last_frame_path], [tracking_points, input_image, input_image_end])
-    
-        run_button.click(Framer.run, [first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id], output_video)
-    
+                """
+            )
+
+        image_upload_button.upload(
+            fn=preprocess_image,
+            inputs=image_upload_button,
+            outputs=[input_image, first_frame_path, tracking_points],
+        )
+
+        image_end_upload_button.upload(
+            fn=preprocess_image_end,
+            inputs=image_end_upload_button,
+            outputs=[input_image_end, last_frame_path, tracking_points],
+        )
+
+        add_drag_button.click(
+            fn=add_drag,
+            inputs=tracking_points,
+            outputs=tracking_points,
+        )
+
+        delete_last_drag_button.click(
+            fn=delete_last_drag,
+            inputs=[tracking_points, first_frame_path, last_frame_path],
+            outputs=[tracking_points, input_image, input_image_end],
+        )
+
+        delete_last_step_button.click(
+            fn=delete_last_step,
+            inputs=[tracking_points, first_frame_path, last_frame_path],
+            outputs=[tracking_points, input_image, input_image_end],
+        )
+
+        reset_button.click(
+            fn=reset_states,
+            outputs=[input_image, input_image_end, first_frame_path, last_frame_path, output_video, tracking_points],
+        )
+
+        gr.on(
+            triggers=[input_image.select, input_image_end.select],
+            fn=add_tracking_points,
+            inputs=[tracking_points, first_frame_path, last_frame_path],
+            outputs=[tracking_points, input_image, input_image_end],
+        )
+
+        run_button.click(
+            fn=run,
+            inputs=[first_frame_path, last_frame_path, tracking_points, controlnet_cond_scale, motion_bucket_id],
+            outputs=output_video,
+        )
+
     demo.launch()

requirements.txt

@@ -1,14 +1,279 @@
-torch==2.0.0
-torchvision
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements.txt
+accelerate==1.1.1
+    # via framer (pyproject.toml)
+aiofiles==23.2.1
+    # via gradio
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.6.2.post1
+    # via
+    #   gradio
+    #   httpx
+    #   starlette
+av==13.1.0
+    # via framer (pyproject.toml)
+certifi==2024.8.30
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+charset-normalizer==3.4.0
+    # via requests
+click==8.1.7
+    # via
+    #   typer
+    #   uvicorn
+colorlog==6.9.0
+    # via framer (pyproject.toml)
 diffusers==0.24.0
-transformers==4.27.0
-xformers==0.0.18
-imageio==2.27.0
-decord==0.6.0
-einops
-opencv-python
-av
-accelerate==0.27.2
-scipy
-colorlog
-numpy==1.24.3
+    # via framer (pyproject.toml)
+einops==0.8.0
+    # via framer (pyproject.toml)
+exceptiongroup==1.2.2
+    # via anyio
+fastapi==0.115.4
+    # via gradio
+ffmpy==0.4.0
+    # via gradio
+filelock==3.16.1
+    # via
+    #   diffusers
+    #   huggingface-hub
+    #   torch
+    #   transformers
+    #   triton
+fsspec==2024.10.0
+    # via
+    #   gradio-client
+    #   huggingface-hub
+    #   torch
+gradio==5.5.0
+    # via
+    #   framer (pyproject.toml)
+    #   spaces
+gradio-client==1.4.2
+    # via gradio
+h11==0.14.0
+    # via
+    #   httpcore
+    #   uvicorn
+hf-transfer==0.1.8
+    # via framer (pyproject.toml)
+httpcore==1.0.6
+    # via httpx
+httpx==0.27.2
+    # via
+    #   gradio
+    #   gradio-client
+    #   safehttpx
+    #   spaces
+huggingface-hub==0.25.2
+    # via
+    #   framer (pyproject.toml)
+    #   accelerate
+    #   diffusers
+    #   gradio
+    #   gradio-client
+    #   tokenizers
+    #   transformers
+idna==3.10
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+imageio==2.36.0
+    # via framer (pyproject.toml)
+importlib-metadata==8.5.0
+    # via diffusers
+jinja2==3.1.4
+    # via
+    #   gradio
+    #   torch
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==2.1.5
+    # via
+    #   gradio
+    #   jinja2
+mdurl==0.1.2
+    # via markdown-it-py
+mpmath==1.3.0
+    # via sympy
+networkx==3.4.2
+    # via torch
+numpy==1.24.3
+    # via
+    #   accelerate
+    #   diffusers
+    #   gradio
+    #   imageio
+    #   opencv-python
+    #   pandas
+    #   scipy
+    #   torchvision
+    #   transformers
+nvidia-cublas-cu12==12.1.3.1
+    # via
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cuda-cupti-cu12==12.1.105
+    # via torch
+nvidia-cuda-nvrtc-cu12==12.1.105
+    # via torch
+nvidia-cuda-runtime-cu12==12.1.105
+    # via torch
+nvidia-cudnn-cu12==9.1.0.70
+    # via torch
+nvidia-cufft-cu12==11.0.2.54
+    # via torch
+nvidia-curand-cu12==10.3.2.106
+    # via torch
+nvidia-cusolver-cu12==11.4.5.107
+    # via torch
+nvidia-cusparse-cu12==12.1.0.106
+    # via
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-nccl-cu12==2.20.5
+    # via torch
+nvidia-nvjitlink-cu12==12.6.77
+    # via
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
+nvidia-nvtx-cu12==12.1.105
+    # via torch
+opencv-python==4.10.0.84
+    # via framer (pyproject.toml)
+orjson==3.10.11
+    # via gradio
+packaging==24.2
+    # via
+    #   accelerate
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   spaces
+    #   transformers
+pandas==2.2.3
+    # via gradio
+pillow==11.0.0
+    # via
+    #   diffusers
+    #   gradio
+    #   imageio
+    #   torchvision
+psutil==5.9.8
+    # via
+    #   accelerate
+    #   spaces
+pydantic==2.9.2
+    # via
+    #   fastapi
+    #   gradio
+    #   spaces
+pydantic-core==2.23.4
+    # via pydantic
+pydub==0.25.1
+    # via gradio
+pygments==2.18.0
+    # via rich
+python-dateutil==2.9.0.post0
+    # via pandas
+python-multipart==0.0.12
+    # via gradio
+pytz==2024.2
+    # via pandas
+pyyaml==6.0.2
+    # via
+    #   accelerate
+    #   gradio
+    #   huggingface-hub
+    #   transformers
+regex==2024.11.6
+    # via
+    #   diffusers
+    #   transformers
+requests==2.32.3
+    # via
+    #   diffusers
+    #   huggingface-hub
+    #   spaces
+    #   transformers
+rich==13.9.4
+    # via typer
+ruff==0.7.3
+    # via gradio
+safehttpx==0.1.1
+    # via gradio
+safetensors==0.4.5
+    # via
+    #   accelerate
+    #   diffusers
+    #   transformers
+scipy==1.14.1
+    # via framer (pyproject.toml)
+semantic-version==2.10.0
+    # via gradio
+shellingham==1.5.4
+    # via typer
+six==1.16.0
+    # via python-dateutil
+sniffio==1.3.1
+    # via
+    #   anyio
+    #   httpx
+spaces==0.30.4
+    # via framer (pyproject.toml)
+starlette==0.41.2
+    # via
+    #   fastapi
+    #   gradio
+sympy==1.13.3
+    # via torch
+tokenizers==0.20.3
+    # via transformers
+tomlkit==0.12.0
+    # via gradio
+torch==2.4.0
+    # via
+    #   framer (pyproject.toml)
+    #   accelerate
+    #   torchvision
+torchvision==0.19.0
+    # via framer (pyproject.toml)
+tqdm==4.67.0
+    # via
+    #   huggingface-hub
+    #   transformers
+transformers==4.46.2
+    # via framer (pyproject.toml)
+triton==3.0.0
+    # via torch
+typer==0.13.0
+    # via gradio
+typing-extensions==4.12.2
+    # via
+    #   anyio
+    #   fastapi
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   pydantic
+    #   pydantic-core
+    #   rich
+    #   spaces
+    #   torch
+    #   typer
+    #   uvicorn
+tzdata==2024.2
+    # via pandas
+urllib3==2.2.3
+    # via requests
+uvicorn==0.32.0
+    # via gradio
+websockets==12.0
+    # via gradio-client
+zipp==3.21.0
+    # via importlib-metadata