| 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/" | |
| VERSION2SPECS = { | |
| "svd": { | |
| "T": 14, | |
| "H": 576, | |
| "W": 1024, | |
| "C": 4, | |
| "f": 8, | |
| "config": "configs/inference/svd.yaml", | |
| "ckpt": "checkpoints/svd.safetensors", | |
| "options": { | |
| "discretization": 1, | |
| "cfg": 2.5, | |
| "sigma_min": 0.002, | |
| "sigma_max": 700.0, | |
| "rho": 7.0, | |
| "guider": 2, | |
| "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"], | |
| "num_steps": 25, | |
| }, | |
| }, | |
| "svd_image_decoder": { | |
| "T": 14, | |
| "H": 576, | |
| "W": 1024, | |
| "C": 4, | |
| "f": 8, | |
| "config": "configs/inference/svd_image_decoder.yaml", | |
| "ckpt": "checkpoints/svd_image_decoder.safetensors", | |
| "options": { | |
| "discretization": 1, | |
| "cfg": 2.5, | |
| "sigma_min": 0.002, | |
| "sigma_max": 700.0, | |
| "rho": 7.0, | |
| "guider": 2, | |
| "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"], | |
| "num_steps": 25, | |
| }, | |
| }, | |
| "svd_xt": { | |
| "T": 25, | |
| "H": 576, | |
| "W": 1024, | |
| "C": 4, | |
| "f": 8, | |
| "config": "configs/inference/svd.yaml", | |
| "ckpt": "checkpoints/svd_xt.safetensors", | |
| "options": { | |
| "discretization": 1, | |
| "cfg": 3.0, | |
| "min_cfg": 1.5, | |
| "sigma_min": 0.002, | |
| "sigma_max": 700.0, | |
| "rho": 7.0, | |
| "guider": 2, | |
| "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"], | |
| "num_steps": 30, | |
| "decoding_t": 14, | |
| }, | |
| }, | |
| "svd_xt_image_decoder": { | |
| "T": 25, | |
| "H": 576, | |
| "W": 1024, | |
| "C": 4, | |
| "f": 8, | |
| "config": "configs/inference/svd_image_decoder.yaml", | |
| "ckpt": "checkpoints/svd_xt_image_decoder.safetensors", | |
| "options": { | |
| "discretization": 1, | |
| "cfg": 3.0, | |
| "min_cfg": 1.5, | |
| "sigma_min": 0.002, | |
| "sigma_max": 700.0, | |
| "rho": 7.0, | |
| "guider": 2, | |
| "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"], | |
| "num_steps": 30, | |
| "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 / SV3D") | |
| version = st.selectbox( | |
| "Model Version", | |
| [k for k in VERSION2SPECS.keys()], | |
| 0, | |
| ) | |
| version_dict = VERSION2SPECS[version] | |
| if st.checkbox("Load Model"): | |
| mode = "img2vid" | |
| else: | |
| mode = "skip" | |
| H = st.sidebar.number_input( | |
| "H", value=version_dict["H"], min_value=64, max_value=2048 | |
| ) | |
| W = st.sidebar.number_input( | |
| "W", value=version_dict["W"], min_value=64, max_value=2048 | |
| ) | |
| T = st.sidebar.number_input( | |
| "T", value=version_dict["T"], min_value=0, max_value=128 | |
| ) | |
| C = version_dict["C"] | |
| F = version_dict["f"] | |
| options = version_dict["options"] | |
| if mode != "skip": | |
| state = init_st(version_dict, load_filter=True) | |
| if state["msg"]: | |
| st.info(state["msg"]) | |
| model = state["model"] | |
| ukeys = set( | |
| get_unique_embedder_keys_from_conditioner(state["model"].conditioner) | |
| ) | |
| value_dict = init_embedder_options( | |
| ukeys, | |
| {}, | |
| ) | |
| 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) | |
| 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) | |
| ) | |
| seed_everything(seed) | |
| save_locally, save_path = init_save_locally( | |
| 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) | |
| num_samples = num_rows * num_cols | |
| decoding_t = st.number_input( | |
| "Decode t frames at a time (set small if you are low on VRAM)", | |
| value=options.get("decoding_t", T), | |
| min_value=1, | |
| max_value=int(1e9), | |
| ) | |
| if st.checkbox("Overwrite fps in mp4 generator", False): | |
| saving_fps = st.number_input( | |
| f"saving video at fps:", value=value_dict["fps"], min_value=1 | |
| ) | |
| else: | |
| saving_fps = value_dict["fps"] | |
| if st.button("Sample"): | |
| out = do_sample( | |
| model, | |
| sampler, | |
| value_dict, | |
| num_samples, | |
| H, | |
| W, | |
| C, | |
| F, | |
| T=T, | |
| batch2model_input=["num_video_frames", "image_only_indicator"], | |
| force_uc_zero_embeddings=options.get("force_uc_zero_embeddings", None), | |
| force_cond_zero_embeddings=options.get( | |
| "force_cond_zero_embeddings", None | |
| ), | |
| return_latents=False, | |
| decoding_t=decoding_t, | |
| ) | |
| if isinstance(out, (tuple, list)): | |
| samples, samples_z = out | |
| else: | |
| samples = out | |
| samples_z = None | |
| if save_locally: | |
| save_video_as_grid_and_mp4(samples, save_path, T, fps=saving_fps) | |