Drop class and use functionality for user-uploaded models

app.py

@@ -28,224 +28,207 @@ import numpy as np
 from scipy.interpolate import PchipInterpolator
 
 SAVE_DIR = "outputs"
+os.makedirs(SAVE_DIR, exist_ok=True)
 LENGTH = 16
 WIDTH = 512
 HEIGHT = 320
 LAYER_CAPACITY = 4
 DEVICE = "cuda"
-
+WEIGHT_DTYPE = torch.bfloat16
+PIPELINE = None
+GENERATOR = None
 os.makedirs("checkpoints", exist_ok=True)
 
 snapshot_download(
     "Yuppie1204/LayerAnimate-Mix",
     local_dir="checkpoints/LayerAnimate-Mix",
 )
-
-class LayerAnimate:
-
-    @spaces.GPU
-    def __init__(self):
-        self.savedir = SAVE_DIR
-        os.makedirs(self.savedir, exist_ok=True)
-
-        self.weight_dtype  = torch.bfloat16
-        self.device        = DEVICE
-        self.text_encoder  = FrozenOpenCLIPEmbedder().eval()
-        self.image_encoder = FrozenOpenCLIPImageEmbedderV2().eval()
-
-        self.W = WIDTH
-        self.H = HEIGHT
-        self.L = LENGTH
-        self.layer_capacity = LAYER_CAPACITY
-
-        self.transforms = transforms.Compose([
-            transforms.Resize(min(self.H, self.W)),
-            transforms.CenterCrop((self.H, self.W)),
-        ])
-        self.pipeline = None
-        self.generator = None
-        # sample_grid is used to generate fixed trajectories to freeze static layers
-        self.sample_grid = np.meshgrid(np.linspace(0, self.W - 1, 10, dtype=int), np.linspace(0, self.H - 1, 10, dtype=int))
-        self.sample_grid = np.stack(self.sample_grid, axis=-1).reshape(-1, 1, 2)
-        self.sample_grid = np.repeat(self.sample_grid, self.L, axis=1) # [N, F, 2]
-        self.set_model("checkpoints/LayerAnimate-Mix")
-
-    @spaces.GPU
-    def set_seed(self, seed):
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        self.generator = torch.Generator(self.device).manual_seed(seed)
-
-    @spaces.GPU
-    def set_model(self, pretrained_model_path):
-        scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
-        image_projector = Resampler.from_pretrained(pretrained_model_path, subfolder="image_projector").eval()
-        vae, vae_dualref = None, None
-        if "I2V" or "Mix" in pretrained_model_path:
-            vae           = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").eval()
-        if "Interp" or "Mix" in pretrained_model_path:
-            vae_dualref   = AutoencoderKL_Dualref.from_pretrained(pretrained_model_path, subfolder="vae_dualref").eval()
-        unet              = UNetModel.from_pretrained(pretrained_model_path, subfolder="unet").eval()
-        layer_controlnet  = LayerControlNet.from_pretrained(pretrained_model_path, subfolder="layer_controlnet").eval()
-
-        self.pipeline = AnimationPipeline(
-            vae=vae, vae_dualref=vae_dualref, text_encoder=self.text_encoder, image_encoder=self.image_encoder, image_projector=image_projector,
-            unet=unet, layer_controlnet=layer_controlnet, scheduler=scheduler
-        ).to(device=self.device, dtype=self.weight_dtype)
-        if "Interp" or "Mix" in pretrained_model_path:
-            self.pipeline.vae_dualref.decoder.to(dtype=torch.float32)
-        return pretrained_model_path
-
-    def upload_image(self, image):
-        image = self.transforms(image)
-        return image
-
-    def run(self, input_image, input_image_end, pretrained_model_path, seed,
-            prompt, n_prompt, num_inference_steps, guidance_scale,
-            *layer_args):
-        self.set_seed(seed)
-        global layer_tracking_points
-        args_layer_tracking_points = [layer_tracking_points[i].value for i in range(self.layer_capacity)]
-
-        args_layer_masks = layer_args[:self.layer_capacity]
-        args_layer_masks_end = layer_args[self.layer_capacity : 2 * self.layer_capacity]
-        args_layer_controls = layer_args[2 * self.layer_capacity : 3 * self.layer_capacity]
-        args_layer_scores = list(layer_args[3 * self.layer_capacity : 4 * self.layer_capacity])
-        args_layer_sketches = layer_args[4 * self.layer_capacity : 5 * self.layer_capacity]
-        args_layer_valids = layer_args[5 * self.layer_capacity : 6 * self.layer_capacity]
-        args_layer_statics = layer_args[6 * self.layer_capacity : 7 * self.layer_capacity]
-        for layer_idx in range(self.layer_capacity):
-            if args_layer_controls[layer_idx] != "score":
-                args_layer_scores[layer_idx] = -1
-            if args_layer_statics[layer_idx]:
-                args_layer_scores[layer_idx] = 0
-
-        mode = "i2v"
-        image1 = F.to_tensor(input_image) * 2 - 1
-        frame_tensor = image1[None].to(self.device) # [F, C, H, W]
-        if input_image_end is not None:
-            mode = "interpolate"
-            image2 = F.to_tensor(input_image_end) * 2 - 1
-            frame_tensor2 = image2[None].to(self.device)
-            frame_tensor = torch.cat([frame_tensor, frame_tensor2], dim=0)
-        frame_tensor = frame_tensor[None]
-
-        if mode == "interpolate":
-            layer_masks = torch.zeros((1, self.layer_capacity, 2, 1, self.H, self.W), dtype=torch.bool)
+TEXT_ENCODER  = FrozenOpenCLIPEmbedder().eval()
+IMAGE_ENCODER = FrozenOpenCLIPImageEmbedderV2().eval()
+
+TRANSFORMS = transforms.Compose([
+    transforms.Resize(min(HEIGHT, WIDTH)),
+    transforms.CenterCrop((HEIGHT, WIDTH)),
+])
+SAMPLE_GRID = np.meshgrid(np.linspace(0, WIDTH - 1, 10, dtype=int), np.linspace(0, HEIGHT - 1, 10, dtype=int))
+SAMPLE_GRID = np.stack(SAMPLE_GRID, axis=-1).reshape(-1, 1, 2)
+SAMPLE_GRID = np.repeat(SAMPLE_GRID, LENGTH, axis=1) # [N, F, 2]
+
+@spaces.GPU
+def set_seed(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    return torch.Generator(DEVICE).manual_seed(seed)
+
+@spaces.GPU
+def set_model(pretrained_model_path):
+    global PIPELINE
+    scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
+    image_projector = Resampler.from_pretrained(pretrained_model_path, subfolder="image_projector").eval()
+    vae, vae_dualref = None, None
+    if "I2V" or "Mix" in pretrained_model_path:
+        vae           = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").eval()
+    if "Interp" or "Mix" in pretrained_model_path:
+        vae_dualref   = AutoencoderKL_Dualref.from_pretrained(pretrained_model_path, subfolder="vae_dualref").eval()
+    unet              = UNetModel.from_pretrained(pretrained_model_path, subfolder="unet").eval()
+    layer_controlnet  = LayerControlNet.from_pretrained(pretrained_model_path, subfolder="layer_controlnet").eval()
+
+    PIPELINE = AnimationPipeline(
+        vae=vae, vae_dualref=vae_dualref, text_encoder=TEXT_ENCODER, image_encoder=IMAGE_ENCODER, image_projector=image_projector,
+        unet=unet, layer_controlnet=layer_controlnet, scheduler=scheduler
+    ).to(device=DEVICE, dtype=WEIGHT_DTYPE)
+    if "Interp" or "Mix" in pretrained_model_path:
+        PIPELINE.vae_dualref.decoder.to(dtype=torch.float32)
+    return pretrained_model_path
+
+set_model("checkpoints/LayerAnimate-Mix")
+
+def upload_image(image):
+    image = TRANSFORMS(image)
+    return image
+
+def run(input_image, input_image_end, pretrained_model_path, seed,
+        prompt, n_prompt, num_inference_steps, guidance_scale,
+        *layer_args):
+    generator = set_seed(seed)
+    global layer_tracking_points
+    args_layer_tracking_points = [layer_tracking_points[i].value for i in range(LAYER_CAPACITY)]
+
+    args_layer_masks = layer_args[:LAYER_CAPACITY]
+    args_layer_masks_end = layer_args[LAYER_CAPACITY : 2 * LAYER_CAPACITY]
+    args_layer_controls = layer_args[2 * LAYER_CAPACITY : 3 * LAYER_CAPACITY]
+    args_layer_scores = list(layer_args[3 * LAYER_CAPACITY : 4 * LAYER_CAPACITY])
+    args_layer_sketches = layer_args[4 * LAYER_CAPACITY : 5 * LAYER_CAPACITY]
+    args_layer_valids = layer_args[5 * LAYER_CAPACITY : 6 * LAYER_CAPACITY]
+    args_layer_statics = layer_args[6 * LAYER_CAPACITY : 7 * LAYER_CAPACITY]
+    for layer_idx in range(LAYER_CAPACITY):
+        if args_layer_controls[layer_idx] != "score":
+            args_layer_scores[layer_idx] = -1
+        if args_layer_statics[layer_idx]:
+            args_layer_scores[layer_idx] = 0
+
+    mode = "i2v"
+    image1 = F.to_tensor(input_image) * 2 - 1
+    frame_tensor = image1[None].to(DEVICE) # [F, C, H, W]
+    if input_image_end is not None:
+        mode = "interpolate"
+        image2 = F.to_tensor(input_image_end) * 2 - 1
+        frame_tensor2 = image2[None].to(DEVICE)
+        frame_tensor = torch.cat([frame_tensor, frame_tensor2], dim=0)
+    frame_tensor = frame_tensor[None]
+
+    if mode == "interpolate":
+        layer_masks = torch.zeros((1, LAYER_CAPACITY, 2, 1, HEIGHT, WIDTH), dtype=torch.bool)
+    else:
+        layer_masks = torch.zeros((1, LAYER_CAPACITY, 1, 1, HEIGHT, WIDTH), dtype=torch.bool)
+    for layer_idx in range(LAYER_CAPACITY):
+        if args_layer_masks[layer_idx] is not None:
+            mask = F.to_tensor(args_layer_masks[layer_idx]) > 0.5
+            layer_masks[0, layer_idx, 0] = mask
+        if args_layer_masks_end[layer_idx] is not None and mode == "interpolate":
+            mask = F.to_tensor(args_layer_masks_end[layer_idx]) > 0.5
+            layer_masks[0, layer_idx, 1] = mask
+    layer_masks = layer_masks.to(DEVICE)
+    layer_regions = layer_masks * frame_tensor[:, None]
+    layer_validity = torch.tensor([args_layer_valids], dtype=torch.bool, device=DEVICE)
+    motion_scores = torch.tensor([args_layer_scores], dtype=WEIGHT_DTYPE, device=DEVICE)
+    layer_static = torch.tensor([args_layer_statics], dtype=torch.bool, device=DEVICE)
+
+    sketch = torch.ones((1, LAYER_CAPACITY, LENGTH, 3, HEIGHT, WIDTH), dtype=WEIGHT_DTYPE)
+    for layer_idx in range(LAYER_CAPACITY):
+        sketch_path = args_layer_sketches[layer_idx]
+        if sketch_path is not None:
+            video_reader = decord.VideoReader(sketch_path)
+            assert len(video_reader) == LENGTH, f"Input the length of sketch sequence should match the video length."
+            video_frames = video_reader.get_batch(range(LENGTH)).asnumpy()
+            sketch_values = [F.to_tensor(TRANSFORMS(Image.fromarray(frame))) for frame in video_frames]
+            sketch_values = torch.stack(sketch_values) * 2 - 1
+            sketch[0, layer_idx] = sketch_values
+    sketch = sketch.to(DEVICE)
+
+    heatmap = torch.zeros((1, LAYER_CAPACITY, LENGTH, 3, HEIGHT, WIDTH), dtype=WEIGHT_DTYPE)
+    heatmap[:, :, :, 0] -= 1
+    trajectory = []
+    traj_layer_index = []
+    for layer_idx in range(LAYER_CAPACITY):
+        tracking_points = args_layer_tracking_points[layer_idx]
+        if args_layer_statics[layer_idx]:
+            # generate pseudo trajectory for static layers
+            temp_layer_mask = layer_masks[0, layer_idx, 0, 0].cpu().numpy()
+            valid_flag = temp_layer_mask[SAMPLE_GRID[:, 0, 1], SAMPLE_GRID[:, 0, 0]]
+            valid_grid = SAMPLE_GRID[valid_flag]    # [F, N, 2]
+            trajectory.extend(list(valid_grid))
+            traj_layer_index.extend([layer_idx] * valid_grid.shape[0])
         else:
-            layer_masks = torch.zeros((1, self.layer_capacity, 1, 1, self.H, self.W), dtype=torch.bool)
-        for layer_idx in range(self.layer_capacity):
-            if args_layer_masks[layer_idx] is not None:
-                mask = F.to_tensor(args_layer_masks[layer_idx]) > 0.5
-                layer_masks[0, layer_idx, 0] = mask
-            if args_layer_masks_end[layer_idx] is not None and mode == "interpolate":
-                mask = F.to_tensor(args_layer_masks_end[layer_idx]) > 0.5
-                layer_masks[0, layer_idx, 1] = mask
-        layer_masks = layer_masks.to(self.device)
-        layer_regions = layer_masks * frame_tensor[:, None]
-        layer_validity = torch.tensor([args_layer_valids], dtype=torch.bool, device=self.device)
-        motion_scores = torch.tensor([args_layer_scores], dtype=self.weight_dtype, device=self.device)
-        layer_static = torch.tensor([args_layer_statics], dtype=torch.bool, device=self.device)
-
-        sketch = torch.ones((1, self.layer_capacity, self.L, 3, self.H, self.W), dtype=self.weight_dtype)
-        for layer_idx in range(self.layer_capacity):
-            sketch_path = args_layer_sketches[layer_idx]
-            if sketch_path is not None:
-                video_reader = decord.VideoReader(sketch_path)
-                assert len(video_reader) == self.L, f"Input the length of sketch sequence should match the video length."
-                video_frames = video_reader.get_batch(range(self.L)).asnumpy()
-                sketch_values = [F.to_tensor(self.transforms(Image.fromarray(frame))) for frame in video_frames]
-                sketch_values = torch.stack(sketch_values) * 2 - 1
-                sketch[0, layer_idx] = sketch_values
-        sketch = sketch.to(self.device)
-
-        heatmap = torch.zeros((1, self.layer_capacity, self.L, 3, self.H, self.W), dtype=self.weight_dtype)
-        heatmap[:, :, :, 0] -= 1
-        trajectory = []
-        traj_layer_index = []
-        for layer_idx in range(self.layer_capacity):
-            tracking_points = args_layer_tracking_points[layer_idx]
-            if args_layer_statics[layer_idx]:
-                # generate pseudo trajectory for static layers
-                temp_layer_mask = layer_masks[0, layer_idx, 0, 0].cpu().numpy()
-                valid_flag = temp_layer_mask[self.sample_grid[:, 0, 1], self.sample_grid[:, 0, 0]]
-                valid_grid = self.sample_grid[valid_flag]    # [F, N, 2]
-                trajectory.extend(list(valid_grid))
-                traj_layer_index.extend([layer_idx] * valid_grid.shape[0])
-            else:
-                for temp_track in tracking_points:
-                    if len(temp_track) > 1:
-                        x = [point[0] for point in temp_track]
-                        y = [point[1] for point in temp_track]
-                        t = np.linspace(0, 1, len(temp_track))
-                        fx = PchipInterpolator(t, x)
-                        fy = PchipInterpolator(t, y)
-                        t_new = np.linspace(0, 1, self.L)
-                        x_new = fx(t_new)
-                        y_new = fy(t_new)
-                        temp_traj = np.stack([x_new, y_new], axis=-1).astype(np.float32)
-                        trajectory.append(temp_traj)
-                        traj_layer_index.append(layer_idx)
-                    elif len(temp_track) == 1:
-                        trajectory.append(np.array(temp_track * self.L))
-                        traj_layer_index.append(layer_idx)
-
-        trajectory = np.stack(trajectory)
-        trajectory = np.transpose(trajectory, (1, 0, 2))
-        traj_layer_index = np.array(traj_layer_index)
-        heatmap = generate_gaussian_heatmap(trajectory, self.W, self.H, traj_layer_index, self.layer_capacity, offset=True)
-        heatmap = rearrange(heatmap, "f n c h w -> (f n) c h w")
-        graymap, offset = heatmap[:, :1], heatmap[:, 1:]
-        graymap = graymap / 255.
-        rad = torch.sqrt(offset[:, 0:1]**2 + offset[:, 1:2]**2)
-        rad_max = torch.max(rad)
-        epsilon = 1e-5
-        offset = offset / (rad_max + epsilon)
-        graymap = graymap * 2 - 1
-        heatmap = torch.cat([graymap, offset], dim=1)
-        heatmap = rearrange(heatmap, '(f n) c h w -> n f c h w', n=self.layer_capacity)
-        heatmap = heatmap[None]
-        heatmap = heatmap.to(self.device)
-
-        sample = self.pipeline(
-            prompt,
-            self.L,
-            self.H,
-            self.W,
-            frame_tensor,
-            layer_masks             = layer_masks,
-            layer_regions           = layer_regions,
-            layer_static            = layer_static,
-            motion_scores           = motion_scores,
-            sketch                  = sketch,
-            trajectory              = heatmap,
-            layer_validity          = layer_validity,
-            num_inference_steps     = num_inference_steps,
-            guidance_scale          = guidance_scale,
-            guidance_rescale        = 0.7,
-            negative_prompt         = n_prompt,
-            num_videos_per_prompt   = 1,
-            eta                     = 1.0,
-            generator               = self.generator,
-            fps                     = 24,
-            mode                    = mode,
-            weight_dtype            = self.weight_dtype,
-            output_type             = "tensor",
-        ).videos
-        output_video_path = os.path.join(self.savedir, "video.mp4")
-        save_videos_grid(sample, output_video_path, fps=8)
-        output_video_traj_path = os.path.join(self.savedir, "video_with_traj.mp4")
-        vis_traj_flag = np.zeros(trajectory.shape[1], dtype=bool)
-        for traj_idx in range(trajectory.shape[1]):
-            if not args_layer_statics[traj_layer_index[traj_idx]]:
-                vis_traj_flag[traj_idx] = True
-        vis_traj = torch.from_numpy(trajectory[:, vis_traj_flag])
-        save_videos_with_traj(sample[0], vis_traj, os.path.join(self.savedir, f"video_with_traj.mp4"), fps=8, line_width=7, circle_radius=10)
-        return output_video_path, output_video_traj_path
-
-
-layeranimate = LayerAnimate()
-
+            for temp_track in tracking_points:
+                if len(temp_track) > 1:
+                    x = [point[0] for point in temp_track]
+                    y = [point[1] for point in temp_track]
+                    t = np.linspace(0, 1, len(temp_track))
+                    fx = PchipInterpolator(t, x)
+                    fy = PchipInterpolator(t, y)
+                    t_new = np.linspace(0, 1, LENGTH)
+                    x_new = fx(t_new)
+                    y_new = fy(t_new)
+                    temp_traj = np.stack([x_new, y_new], axis=-1).astype(np.float32)
+                    trajectory.append(temp_traj)
+                    traj_layer_index.append(layer_idx)
+                elif len(temp_track) == 1:
+                    trajectory.append(np.array(temp_track * LENGTH))
+                    traj_layer_index.append(layer_idx)
+
+    trajectory = np.stack(trajectory)
+    trajectory = np.transpose(trajectory, (1, 0, 2))
+    traj_layer_index = np.array(traj_layer_index)
+    heatmap = generate_gaussian_heatmap(trajectory, WIDTH, HEIGHT, traj_layer_index, LAYER_CAPACITY, offset=True)
+    heatmap = rearrange(heatmap, "f n c h w -> (f n) c h w")
+    graymap, offset = heatmap[:, :1], heatmap[:, 1:]
+    graymap = graymap / 255.
+    rad = torch.sqrt(offset[:, 0:1]**2 + offset[:, 1:2]**2)
+    rad_max = torch.max(rad)
+    epsilon = 1e-5
+    offset = offset / (rad_max + epsilon)
+    graymap = graymap * 2 - 1
+    heatmap = torch.cat([graymap, offset], dim=1)
+    heatmap = rearrange(heatmap, '(f n) c h w -> n f c h w', n=LAYER_CAPACITY)
+    heatmap = heatmap[None]
+    heatmap = heatmap.to(DEVICE)
+
+    sample = PIPELINE(
+        prompt,
+        LENGTH,
+        HEIGHT,
+        WIDTH,
+        frame_tensor,
+        layer_masks             = layer_masks,
+        layer_regions           = layer_regions,
+        layer_static            = layer_static,
+        motion_scores           = motion_scores,
+        sketch                  = sketch,
+        trajectory              = heatmap,
+        layer_validity          = layer_validity,
+        num_inference_steps     = num_inference_steps,
+        guidance_scale          = guidance_scale,
+        guidance_rescale        = 0.7,
+        negative_prompt         = n_prompt,
+        num_videos_per_prompt   = 1,
+        eta                     = 1.0,
+        generator               = generator,
+        fps                     = 24,
+        mode                    = mode,
+        weight_dtype            = WEIGHT_DTYPE,
+        output_type             = "tensor",
+    ).videos
+    output_video_path = os.path.join(SAVE_DIR, "video.mp4")
+    save_videos_grid(sample, output_video_path, fps=8)
+    output_video_traj_path = os.path.join(SAVE_DIR, "video_with_traj.mp4")
+    vis_traj_flag = np.zeros(trajectory.shape[1], dtype=bool)
+    for traj_idx in range(trajectory.shape[1]):
+        if not args_layer_statics[traj_layer_index[traj_idx]]:
+            vis_traj_flag[traj_idx] = True
+    vis_traj = torch.from_numpy(trajectory[:, vis_traj_flag])
+    save_videos_with_traj(sample[0], vis_traj, os.path.join(SAVE_DIR, f"video_with_traj.mp4"), fps=8, line_width=7, circle_radius=10)
+    return output_video_path, output_video_traj_path
 
 def update_layer_region(image, layer_mask):
     if image is None or layer_mask is None:
@@ -558,13 +541,13 @@ if __name__ == "__main__":
                 ```
                 """)
 
-        pretrained_model_path.input(layeranimate.set_model, pretrained_model_path, pretrained_model_path)
-        input_image.upload(layeranimate.upload_image, input_image, input_image)
-        input_image_end.upload(layeranimate.upload_image, input_image_end, input_image_end)
+        pretrained_model_path.input(set_model, pretrained_model_path, pretrained_model_path)
+        input_image.upload(upload_image, input_image, input_image)
+        input_image_end.upload(upload_image, input_image_end, input_image_end)
         for i in range(LAYER_CAPACITY):
-            layer_masks[i].upload(layeranimate.upload_image, layer_masks[i], layer_masks[i])
+            layer_masks[i].upload(upload_image, layer_masks[i], layer_masks[i])
             layer_masks[i].change(update_layer_region, [input_image, layer_masks[i]], [layer_regions[i], layer_valids[i]])
-            layer_masks_end[i].upload(layeranimate.upload_image, layer_masks_end[i], layer_masks_end[i])
+            layer_masks_end[i].upload(upload_image, layer_masks_end[i], layer_masks_end[i])
             layer_masks_end[i].change(update_layer_region, [input_image_end, layer_masks_end[i]], [layer_regions_end[i], layer_valids[i]])
             layer_traj_controls[i][0].click(add_drag, layer_indices[i], None)
             layer_traj_controls[i][1].click(
@@ -598,7 +581,7 @@ if __name__ == "__main__":
                 [layer_regions[i], layer_regions_end[i]]
             )
         run_button.click(
-            layeranimate.run,
+            run,
             [input_image, input_image_end, pretrained_model_path, seed, text_prompt, text_n_prompt, num_inference_steps, guidance_scale,
              *layer_masks, *layer_masks_end, *layer_controls, *layer_score_controls, *layer_sketch_controls, *layer_valids, *layer_statics],
             [output_video, output_video_traj]