SV3D inference code (#300)

* Makes init changes for SV3D

* Small fixes : cond_aug

* Fixes SV3D checkpoint, fixes rembg

* Black formatting

* Adds streamlit demo, fixes simple sample script

* Removes SV3D video_decoder, keeps SV3D image_decoder

* Updates README

* Minor updates

* Remove GSO script

---------

Co-authored-by: Vikram Voleti <vikram@ip-26-0-153-234.us-west-2.compute.internal>

scripts/demo/gradio_app.py

@@ -23,9 +23,11 @@ from PIL import Image
 from torchvision.transforms import ToTensor
 
 from scripts.sampling.simple_video_sample import (
-    get_batch, get_unique_embedder_keys_from_conditioner, load_model)
-from scripts.util.detection.nsfw_and_watermark_dectection import \
-    DeepFloydDataFiltering
+    get_batch,
+    get_unique_embedder_keys_from_conditioner,
+    load_model,
+)
+from scripts.util.detection.nsfw_and_watermark_dectection import DeepFloydDataFiltering
 from sgm.inference.helpers import embed_watermark
 from sgm.util import default, instantiate_from_config
 

scripts/demo/streamlit_helpers.py

@@ -5,6 +5,7 @@ from glob import glob
 from typing import Dict, List, Optional, Tuple, Union
 
 import cv2
+import imageio
 import numpy as np
 import streamlit as st
 import torch
@@ -15,25 +16,30 @@ from imwatermark import WatermarkEncoder
 from omegaconf import ListConfig, OmegaConf
 from PIL import Image
 from safetensors.torch import load_file as load_safetensors
+from scripts.demo.discretization import (
+    Img2ImgDiscretizationWrapper,
+    Txt2NoisyDiscretizationWrapper,
+)
+from scripts.util.detection.nsfw_and_watermark_dectection import DeepFloydDataFiltering
+from sgm.inference.helpers import embed_watermark
+from sgm.modules.diffusionmodules.guiders import (
+    LinearPredictionGuider,
+    TrianglePredictionGuider,
+    VanillaCFG,
+)
+from sgm.modules.diffusionmodules.sampling import (
+    DPMPP2MSampler,
+    DPMPP2SAncestralSampler,
+    EulerAncestralSampler,
+    EulerEDMSampler,
+    HeunEDMSampler,
+    LinearMultistepSampler,
+)
+from sgm.util import append_dims, default, instantiate_from_config
 from torch import autocast
 from torchvision import transforms
 from torchvision.utils import make_grid, save_image
 
-from scripts.demo.discretization import (Img2ImgDiscretizationWrapper,
-                                         Txt2NoisyDiscretizationWrapper)
-from scripts.util.detection.nsfw_and_watermark_dectection import \
-    DeepFloydDataFiltering
-from sgm.inference.helpers import embed_watermark
-from sgm.modules.diffusionmodules.guiders import (LinearPredictionGuider,
-                                                  VanillaCFG)
-from sgm.modules.diffusionmodules.sampling import (DPMPP2MSampler,
-                                                   DPMPP2SAncestralSampler,
-                                                   EulerAncestralSampler,
-                                                   EulerEDMSampler,
-                                                   HeunEDMSampler,
-                                                   LinearMultistepSampler)
-from sgm.util import append_dims, default, instantiate_from_config
-
 
 @st.cache_resource()
 def init_st(version_dict, load_ckpt=True, load_filter=True):
@@ -222,6 +228,7 @@ def get_guider(options, key):
             "VanillaCFG",
             "IdentityGuider",
             "LinearPredictionGuider",
+            "TrianglePredictionGuider",
         ],
         options.get("guider", 0),
     )
@@ -252,7 +259,7 @@ def get_guider(options, key):
             value=options.get("cfg", 1.5),
             min_value=1.0,
         )
-        min_scale = st.number_input(
+        min_scale = st.sidebar.number_input(
             f"min guidance scale",
             value=options.get("min_cfg", 1.0),
             min_value=1.0,
@@ -268,6 +275,29 @@ def get_guider(options, key):
                 **additional_guider_kwargs,
             },
         }
+    elif guider == "TrianglePredictionGuider":
+        max_scale = st.number_input(
+            f"max-cfg-scale #{key}",
+            value=options.get("cfg", 2.5),
+            min_value=1.0,
+            max_value=10.0,
+        )
+        min_scale = st.sidebar.number_input(
+            f"min guidance scale",
+            value=options.get("min_cfg", 1.0),
+            min_value=1.0,
+            max_value=10.0,
+        )
+
+        guider_config = {
+            "target": "sgm.modules.diffusionmodules.guiders.TrianglePredictionGuider",
+            "params": {
+                "max_scale": max_scale,
+                "min_scale": min_scale,
+                "num_frames": options["num_frames"],
+                **additional_guider_kwargs,
+            },
+        }
     else:
         raise NotImplementedError
     return guider_config
@@ -288,8 +318,8 @@ def init_sampling(
             f"num cols #{key}", value=num_cols, min_value=1, max_value=10
         )
 
-    steps = st.sidebar.number_input(
-        f"steps #{key}", value=options.get("num_steps", 40), min_value=1, max_value=1000
+    steps = st.number_input(
+        f"steps #{key}", value=options.get("num_steps", 50), min_value=1, max_value=1000
     )
     sampler = st.sidebar.selectbox(
         f"Sampler #{key}",
@@ -337,13 +367,13 @@ def get_discretization(discretization, options, key=1):
             "target": "sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization",
         }
     elif discretization == "EDMDiscretization":
-        sigma_min = st.number_input(
+        sigma_min = st.sidebar.number_input(
             f"sigma_min #{key}", value=options.get("sigma_min", 0.03)
         )  # 0.0292
-        sigma_max = st.number_input(
+        sigma_max = st.sidebar.number_input(
             f"sigma_max #{key}", value=options.get("sigma_max", 14.61)
         )  # 14.6146
-        rho = st.number_input(f"rho #{key}", value=options.get("rho", 3.0))
+        rho = st.sidebar.number_input(f"rho #{key}", value=options.get("rho", 3.0))
         discretization_config = {
             "target": "sgm.modules.diffusionmodules.discretizer.EDMDiscretization",
             "params": {
@@ -542,7 +572,12 @@ def do_sample(
                         assert T is not None
 
                         if isinstance(
-                            sampler.guider, (VanillaCFG, LinearPredictionGuider)
+                            sampler.guider,
+                            (
+                                VanillaCFG,
+                                LinearPredictionGuider,
+                                TrianglePredictionGuider,
+                            ),
                         ):
                             additional_model_inputs[k] = torch.zeros(
                                 num_samples[0] * 2, num_samples[1]
@@ -678,6 +713,12 @@ def get_batch(
             batch[key] = repeat(
                 value_dict["cond_frames_without_noise"], "1 ... -> b ...", b=N[0]
             )
+        elif key == "polars_rad":
+            batch[key] = torch.tensor(value_dict["polars_rad"]).to(device).repeat(N[0])
+        elif key == "azimuths_rad":
+            batch[key] = (
+                torch.tensor(value_dict["azimuths_rad"]).to(device).repeat(N[0])
+            )
         else:
             batch[key] = value_dict[key]
 
@@ -827,8 +868,13 @@ def load_img_for_prediction(
         st.image(image)
     w, h = image.size
 
-    image = np.array(image).transpose(2, 0, 1)
-    image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
+    image = np.array(image).astype(np.float32) / 255
+    if image.shape[-1] == 4:
+        rgb, alpha = image[:, :, :3], image[:, :, 3:]
+        image = rgb * alpha + (1 - alpha)
+
+    image = image.transpose(2, 0, 1)
+    image = torch.from_numpy(image).to(dtype=torch.float32)
     image = image.unsqueeze(0)
 
     rfs = get_resizing_factor((H, W), (h, w))
@@ -860,28 +906,16 @@ def save_video_as_grid_and_mp4(
         save_image(vid, fp=os.path.join(save_path, f"{base_count:06d}.png"), nrow=4)
 
         video_path = os.path.join(save_path, f"{base_count:06d}.mp4")
-
-        writer = cv2.VideoWriter(
-            video_path,
-            cv2.VideoWriter_fourcc(*"MP4V"),
-            fps,
-            (vid.shape[-1], vid.shape[-2]),
-        )
-
         vid = (
             (rearrange(vid, "t c h w -> t h w c") * 255).cpu().numpy().astype(np.uint8)
         )
-        for frame in vid:
-            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-            writer.write(frame)
-
-        writer.release()
+        imageio.mimwrite(video_path, vid, fps=fps)
 
         video_path_h264 = video_path[:-4] + "_h264.mp4"
-        os.system(f"ffmpeg -i {video_path} -c:v libx264 {video_path_h264}")
-
+        os.system(f"ffmpeg -i '{video_path}' -c:v libx264 '{video_path_h264}'")
         with open(video_path_h264, "rb") as f:
             video_bytes = f.read()
+        os.remove(video_path_h264)
         st.video(video_bytes)
 
         base_count += 1

scripts/demo/sv3d_helpers.py

@@ -0,0 +1,104 @@
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def generate_dynamic_cycle_xy_values(
+    length=21,
+    init_elev=0,
+    num_components=84,
+    frequency_range=(1, 5),
+    amplitude_range=(0.5, 10),
+    step_range=(0, 2),
+):
+    # Y values generation
+    y_sequence = np.ones(length) * init_elev
+    for _ in range(num_components):
+        # Choose a frequency that will complete whole cycles in the sequence
+        frequency = np.random.randint(*frequency_range) * (2 * np.pi / length)
+        amplitude = np.random.uniform(*amplitude_range)
+        phase_shift = np.random.choice([0, np.pi])  # np.random.uniform(0, 2 * np.pi)
+        angles = (
+            np.linspace(0, frequency * length, length, endpoint=False) + phase_shift
+        )
+        y_sequence += np.sin(angles) * amplitude
+    # X values generation
+    # Generate length - 1 steps since the last step is back to start
+    steps = np.random.uniform(*step_range, length - 1)
+    total_step_sum = np.sum(steps)
+    # Calculate the scale factor to scale total steps to just under 360
+    scale_factor = (
+        360 - ((360 / length) * np.random.uniform(*step_range))
+    ) / total_step_sum
+    # Apply the scale factor and generate the sequence of X values
+    x_values = np.cumsum(steps * scale_factor)
+    # Ensure the sequence starts at 0 and add the final step to complete the loop
+    x_values = np.insert(x_values, 0, 0)
+    return x_values, y_sequence
+
+
+def smooth_data(data, window_size):
+    # Extend data at both ends by wrapping around to create a continuous loop
+    pad_size = window_size
+    padded_data = np.concatenate((data[-pad_size:], data, data[:pad_size]))
+
+    # Apply smoothing
+    kernel = np.ones(window_size) / window_size
+    smoothed_data = np.convolve(padded_data, kernel, mode="same")
+
+    # Extract the smoothed data corresponding to the original sequence
+    # Adjust the indices to account for the larger padding
+    start_index = pad_size
+    end_index = -pad_size if pad_size != 0 else None
+    smoothed_original_data = smoothed_data[start_index:end_index]
+    return smoothed_original_data
+
+
+# Function to generate and process the data
+def gen_dynamic_loop(length=21, elev_deg=0):
+    while True:
+        # Generate the combined X and Y values using the new function
+        azim_values, elev_values = generate_dynamic_cycle_xy_values(
+            length=84, init_elev=elev_deg
+        )
+        # Smooth the Y values directly
+        smoothed_elev_values = smooth_data(elev_values, 5)
+        max_magnitude = np.max(np.abs(smoothed_elev_values))
+        if max_magnitude < 90:
+            break
+    subsample = 84 // length
+    azim_rad = np.deg2rad(azim_values[::subsample])
+    elev_rad = np.deg2rad(smoothed_elev_values[::subsample])
+    # Make cond frame the last one
+    return np.roll(azim_rad, -1), np.roll(elev_rad, -1)
+
+
+def plot_3D(azim, polar, save_path, dynamic=True):
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    elev = np.deg2rad(90) - polar
+    fig = plt.figure(figsize=(5, 5))
+    ax = fig.add_subplot(projection="3d")
+    cm = plt.get_cmap("Greys")
+    col_line = [cm(i) for i in np.linspace(0.3, 1, len(azim) + 1)]
+    cm = plt.get_cmap("cool")
+    col = [cm(float(i) / (len(azim))) for i in np.arange(len(azim))]
+    xs = np.cos(elev) * np.cos(azim)
+    ys = np.cos(elev) * np.sin(azim)
+    zs = np.sin(elev)
+    ax.scatter(xs[0], ys[0], zs[0], s=100, color=col[0])
+    xs_d, ys_d, zs_d = (xs[1:] - xs[:-1]), (ys[1:] - ys[:-1]), (zs[1:] - zs[:-1])
+    for i in range(len(xs) - 1):
+        if dynamic:
+            ax.quiver(
+                xs[i], ys[i], zs[i], xs_d[i], ys_d[i], zs_d[i], lw=2, color=col_line[i]
+            )
+        else:
+            ax.plot(xs[i : i + 2], ys[i : i + 2], zs[i : i + 2], lw=2, c=col_line[i])
+        ax.scatter(xs[i + 1], ys[i + 1], zs[i + 1], s=100, color=col[i + 1])
+    ax.scatter(xs[:1], ys[:1], zs[:1], s=120, facecolors="none", edgecolors="k")
+    ax.scatter(xs[-1:], ys[-1:], zs[-1:], s=120, facecolors="none", edgecolors="k")
+    ax.view_init(elev=30, azim=-20, roll=0)
+    plt.savefig(save_path, bbox_inches="tight")
+    plt.clf()
+    plt.close()

scripts/demo/video_sampling.py

@@ -1,8 +1,10 @@
 import os
+import sys
 
+sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), "../../")))
 from pytorch_lightning import seed_everything
-
 from scripts.demo.streamlit_helpers import *
+from scripts.demo.sv3d_helpers import *
 
 SAVE_PATH = "outputs/demo/vid/"
 
@@ -87,11 +89,51 @@ VERSION2SPECS = {
             "decoding_t": 14,
         },
     },
+    "sv3d_u": {
+        "T": 21,
+        "H": 576,
+        "W": 576,
+        "C": 4,
+        "f": 8,
+        "config": "configs/inference/sv3d_u.yaml",
+        "ckpt": "checkpoints/sv3d_u.safetensors",
+        "options": {
+            "discretization": 1,
+            "cfg": 2.5,
+            "sigma_min": 0.002,
+            "sigma_max": 700.0,
+            "rho": 7.0,
+            "guider": 3,
+            "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"],
+            "num_steps": 50,
+            "decoding_t": 14,
+        },
+    },
+    "sv3d_p": {
+        "T": 21,
+        "H": 576,
+        "W": 576,
+        "C": 4,
+        "f": 8,
+        "config": "configs/inference/sv3d_p.yaml",
+        "ckpt": "checkpoints/sv3d_p.safetensors",
+        "options": {
+            "discretization": 1,
+            "cfg": 2.5,
+            "sigma_min": 0.002,
+            "sigma_max": 700.0,
+            "rho": 7.0,
+            "guider": 3,
+            "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"],
+            "num_steps": 50,
+            "decoding_t": 14,
+        },
+    },
 }
 
 
 if __name__ == "__main__":
-    st.title("Stable Video Diffusion")
+    st.title("Stable Video Diffusion / SV3D")
     version = st.selectbox(
         "Model Version",
         [k for k in VERSION2SPECS.keys()],
@@ -131,17 +173,42 @@ if __name__ == "__main__":
             {},
         )
 
+        if "fps" not in ukeys:
+            value_dict["fps"] = 10
+
         value_dict["image_only_indicator"] = 0
 
         if mode == "img2vid":
             img = load_img_for_prediction(W, H)
-            cond_aug = st.number_input(
-                "Conditioning augmentation:", value=0.02, min_value=0.0
-            )
+            if "sv3d" in version:
+                cond_aug = 1e-5
+            else:
+                cond_aug = st.number_input(
+                    "Conditioning augmentation:", value=0.02, min_value=0.0
+                )
             value_dict["cond_frames_without_noise"] = img
             value_dict["cond_frames"] = img + cond_aug * torch.randn_like(img)
             value_dict["cond_aug"] = cond_aug
 
+        if "sv3d_p" in version:
+            elev_deg = st.number_input("elev_deg", value=5, min_value=-90, max_value=90)
+            trajectory = st.selectbox(
+                "Trajectory",
+                ["same elevation", "dynamic"],
+                0,
+            )
+            if trajectory == "same elevation":
+                value_dict["polars_rad"] = np.array([np.deg2rad(90 - elev_deg)] * T)
+                value_dict["azimuths_rad"] = np.linspace(0, 2 * np.pi, T + 1)[1:]
+            elif trajectory == "dynamic":
+                azim_rad, elev_rad = gen_dynamic_loop(length=21, elev_deg=elev_deg)
+                value_dict["polars_rad"] = np.deg2rad(90) - elev_rad
+                value_dict["azimuths_rad"] = azim_rad
+        elif "sv3d_u" in version:
+            elev_deg = st.number_input("elev_deg", value=5, min_value=-90, max_value=90)
+            value_dict["polars_rad"] = np.array([np.deg2rad(90 - elev_deg)] * T)
+            value_dict["azimuths_rad"] = np.linspace(0, 2 * np.pi, T + 1)[1:]
+
         seed = st.sidebar.number_input(
             "seed", value=23, min_value=0, max_value=int(1e9)
         )
@@ -151,6 +218,19 @@ if __name__ == "__main__":
             os.path.join(SAVE_PATH, version), init_value=True
         )
 
+        if "sv3d" in version:
+            plot_save_path = os.path.join(save_path, "plot_3D.png")
+            plot_3D(
+                azim=value_dict["azimuths_rad"],
+                polar=value_dict["polars_rad"],
+                save_path=plot_save_path,
+                dynamic=("sv3d_p" in version),
+            )
+            st.image(
+                plot_save_path,
+                f"3D camera trajectory",
+            )
+
         options["num_frames"] = T
 
         sampler, num_rows, num_cols = init_sampling(options=options)