quantumcontrol/stable-video-diffusion

1from pytorch_lightning import seed_everything2 3from scripts.demo.streamlit_helpers import *4 5SAVE_PATH = "outputs/demo/txt2img/"6 7SD_XL_BASE_RATIOS = {8    "0.5": (704, 1408),9    "0.52": (704, 1344),10    "0.57": (768, 1344),11    "0.6": (768, 1280),12    "0.68": (832, 1216),13    "0.72": (832, 1152),14    "0.78": (896, 1152),15    "0.82": (896, 1088),16    "0.88": (960, 1088),17    "0.94": (960, 1024),18    "1.0": (1024, 1024),19    "1.07": (1024, 960),20    "1.13": (1088, 960),21    "1.21": (1088, 896),22    "1.29": (1152, 896),23    "1.38": (1152, 832),24    "1.46": (1216, 832),25    "1.67": (1280, 768),26    "1.75": (1344, 768),27    "1.91": (1344, 704),28    "2.0": (1408, 704),29    "2.09": (1472, 704),30    "2.4": (1536, 640),31    "2.5": (1600, 640),32    "2.89": (1664, 576),33    "3.0": (1728, 576),34}35 36VERSION2SPECS = {37    "SDXL-base-1.0": {38        "H": 1024,39        "W": 1024,40        "C": 4,41        "f": 8,42        "is_legacy": False,43        "config": "configs/inference/sd_xl_base.yaml",44        "ckpt": "checkpoints/sd_xl_base_1.0.safetensors",45    },46    "SDXL-base-0.9": {47        "H": 1024,48        "W": 1024,49        "C": 4,50        "f": 8,51        "is_legacy": False,52        "config": "configs/inference/sd_xl_base.yaml",53        "ckpt": "checkpoints/sd_xl_base_0.9.safetensors",54    },55    "SD-2.1": {56        "H": 512,57        "W": 512,58        "C": 4,59        "f": 8,60        "is_legacy": True,61        "config": "configs/inference/sd_2_1.yaml",62        "ckpt": "checkpoints/v2-1_512-ema-pruned.safetensors",63    },64    "SD-2.1-768": {65        "H": 768,66        "W": 768,67        "C": 4,68        "f": 8,69        "is_legacy": True,70        "config": "configs/inference/sd_2_1_768.yaml",71        "ckpt": "checkpoints/v2-1_768-ema-pruned.safetensors",72    },73    "SDXL-refiner-0.9": {74        "H": 1024,75        "W": 1024,76        "C": 4,77        "f": 8,78        "is_legacy": True,79        "config": "configs/inference/sd_xl_refiner.yaml",80        "ckpt": "checkpoints/sd_xl_refiner_0.9.safetensors",81    },82    "SDXL-refiner-1.0": {83        "H": 1024,84        "W": 1024,85        "C": 4,86        "f": 8,87        "is_legacy": True,88        "config": "configs/inference/sd_xl_refiner.yaml",89        "ckpt": "checkpoints/sd_xl_refiner_1.0.safetensors",90    },91}92 93 94def load_img(display=True, key=None, device="cuda"):95    image = get_interactive_image(key=key)96    if image is None:97        return None98    if display:99        st.image(image)100    w, h = image.size101    print(f"loaded input image of size ({w}, {h})")102    width, height = map(103        lambda x: x - x % 64, (w, h)104    )  # resize to integer multiple of 64105    image = image.resize((width, height))106    image = np.array(image.convert("RGB"))107    image = image[None].transpose(0, 3, 1, 2)108    image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0109    return image.to(device)110 111 112def run_txt2img(113    state,114    version,115    version_dict,116    is_legacy=False,117    return_latents=False,118    filter=None,119    stage2strength=None,120):121    if version.startswith("SDXL-base"):122        W, H = st.selectbox("Resolution:", list(SD_XL_BASE_RATIOS.values()), 10)123    else:124        H = st.number_input("H", value=version_dict["H"], min_value=64, max_value=2048)125        W = st.number_input("W", value=version_dict["W"], min_value=64, max_value=2048)126    C = version_dict["C"]127    F = version_dict["f"]128 129    init_dict = {130        "orig_width": W,131        "orig_height": H,132        "target_width": W,133        "target_height": H,134    }135    value_dict = init_embedder_options(136        get_unique_embedder_keys_from_conditioner(state["model"].conditioner),137        init_dict,138        prompt=prompt,139        negative_prompt=negative_prompt,140    )141    sampler, num_rows, num_cols = init_sampling(stage2strength=stage2strength)142    num_samples = num_rows * num_cols143 144    if st.button("Sample"):145        st.write(f"**Model I:** {version}")146        out = do_sample(147            state["model"],148            sampler,149            value_dict,150            num_samples,151            H,152            W,153            C,154            F,155            force_uc_zero_embeddings=["txt"] if not is_legacy else [],156            return_latents=return_latents,157            filter=filter,158        )159        return out160 161 162def run_img2img(163    state,164    version_dict,165    is_legacy=False,166    return_latents=False,167    filter=None,168    stage2strength=None,169):170    img = load_img()171    if img is None:172        return None173    H, W = img.shape[2], img.shape[3]174 175    init_dict = {176        "orig_width": W,177        "orig_height": H,178        "target_width": W,179        "target_height": H,180    }181    value_dict = init_embedder_options(182        get_unique_embedder_keys_from_conditioner(state["model"].conditioner),183        init_dict,184        prompt=prompt,185        negative_prompt=negative_prompt,186    )187    strength = st.number_input(188        "**Img2Img Strength**", value=0.75, min_value=0.0, max_value=1.0189    )190    sampler, num_rows, num_cols = init_sampling(191        img2img_strength=strength,192        stage2strength=stage2strength,193    )194    num_samples = num_rows * num_cols195 196    if st.button("Sample"):197        out = do_img2img(198            repeat(img, "1 ... -> n ...", n=num_samples),199            state["model"],200            sampler,201            value_dict,202            num_samples,203            force_uc_zero_embeddings=["txt"] if not is_legacy else [],204            return_latents=return_latents,205            filter=filter,206        )207        return out208 209 210def apply_refiner(211    input,212    state,213    sampler,214    num_samples,215    prompt,216    negative_prompt,217    filter=None,218    finish_denoising=False,219):220    init_dict = {221        "orig_width": input.shape[3] * 8,222        "orig_height": input.shape[2] * 8,223        "target_width": input.shape[3] * 8,224        "target_height": input.shape[2] * 8,225    }226 227    value_dict = init_dict228    value_dict["prompt"] = prompt229    value_dict["negative_prompt"] = negative_prompt230 231    value_dict["crop_coords_top"] = 0232    value_dict["crop_coords_left"] = 0233 234    value_dict["aesthetic_score"] = 6.0235    value_dict["negative_aesthetic_score"] = 2.5236 237    st.warning(f"refiner input shape: {input.shape}")238    samples = do_img2img(239        input,240        state["model"],241        sampler,242        value_dict,243        num_samples,244        skip_encode=True,245        filter=filter,246        add_noise=not finish_denoising,247    )248 249    return samples250 251 252if __name__ == "__main__":253    st.title("Stable Diffusion")254    version = st.selectbox("Model Version", list(VERSION2SPECS.keys()), 0)255    version_dict = VERSION2SPECS[version]256    if st.checkbox("Load Model"):257        mode = st.radio("Mode", ("txt2img", "img2img"), 0)258    else:259        mode = "skip"260    st.write("__________________________")261 262    set_lowvram_mode(st.checkbox("Low vram mode", True))263 264    if version.startswith("SDXL-base"):265        add_pipeline = st.checkbox("Load SDXL-refiner?", False)266        st.write("__________________________")267    else:268        add_pipeline = False269 270    seed = st.sidebar.number_input("seed", value=42, min_value=0, max_value=int(1e9))271    seed_everything(seed)272 273    save_locally, save_path = init_save_locally(os.path.join(SAVE_PATH, version))274 275    if mode != "skip":276        state = init_st(version_dict, load_filter=True)277        if state["msg"]:278            st.info(state["msg"])279        model = state["model"]280 281    is_legacy = version_dict["is_legacy"]282 283    prompt = st.text_input(284        "prompt",285        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",286    )287    if is_legacy:288        negative_prompt = st.text_input("negative prompt", "")289    else:290        negative_prompt = ""  # which is unused291 292    stage2strength = None293    finish_denoising = False294 295    if add_pipeline:296        st.write("__________________________")297        version2 = st.selectbox("Refiner:", ["SDXL-refiner-1.0", "SDXL-refiner-0.9"])298        st.warning(299            f"Running with {version2} as the second stage model. Make sure to provide (V)RAM :) "300        )301        st.write("**Refiner Options:**")302 303        version_dict2 = VERSION2SPECS[version2]304        state2 = init_st(version_dict2, load_filter=False)305        st.info(state2["msg"])306 307        stage2strength = st.number_input(308            "**Refinement strength**", value=0.15, min_value=0.0, max_value=1.0309        )310 311        sampler2, *_ = init_sampling(312            key=2,313            img2img_strength=stage2strength,314            specify_num_samples=False,315        )316        st.write("__________________________")317        finish_denoising = st.checkbox("Finish denoising with refiner.", True)318        if not finish_denoising:319            stage2strength = None320 321    if mode == "txt2img":322        out = run_txt2img(323            state,324            version,325            version_dict,326            is_legacy=is_legacy,327            return_latents=add_pipeline,328            filter=state.get("filter"),329            stage2strength=stage2strength,330        )331    elif mode == "img2img":332        out = run_img2img(333            state,334            version_dict,335            is_legacy=is_legacy,336            return_latents=add_pipeline,337            filter=state.get("filter"),338            stage2strength=stage2strength,339        )340    elif mode == "skip":341        out = None342    else:343        raise ValueError(f"unknown mode {mode}")344    if isinstance(out, (tuple, list)):345        samples, samples_z = out346    else:347        samples = out348        samples_z = None349 350    if add_pipeline and samples_z is not None:351        st.write("**Running Refinement Stage**")352        samples = apply_refiner(353            samples_z,354            state2,355            sampler2,356            samples_z.shape[0],357            prompt=prompt,358            negative_prompt=negative_prompt if is_legacy else "",359            filter=state.get("filter"),360            finish_denoising=finish_denoising,361        )362 363    if save_locally and samples is not None:364        perform_save_locally(save_path, samples)365