diff --git a/README.md b/README.md index 7bc074b..e3ba5c4 100644 --- a/README.md +++ b/README.md @@ -25,12 +25,10 @@ python install.py ## Example Workflow -![example_workflow](https://github.com/user-attachments/assets/afbf7b53-d712-4c9c-a538-53f0dc001f45) +![example_workflow](https://github.com/user-attachments/assets/610d2ccb-c827-4797-9d8f-50b193d43db3) ## Usage -**All the necessary models should be automatically downloaded when the TangoFluxLoader node is used for the first time.** - **Models can also be downloaded using the `install.py` script** ![models_folder_structure](https://github.com/user-attachments/assets/94d8a54a-10d6-4f90-bb4d-3ee181dee3a2) diff --git a/example_workflow.json b/example_workflow.json index aceed95..6adf6c1 100644 --- a/example_workflow.json +++ b/example_workflow.json @@ -1,16 +1,53 @@ { - "last_node_id": 13, - "last_link_id": 15, + "last_node_id": 3, + "last_link_id": 3, "nodes": [ { - "id": 10, - "type": "TangoFluxLoader", + "id": 3, + "type": "TangoFluxVAEDecodeAndPlay", "pos": [ - 380, - 320 + 880, + 230 ], "size": [ - 210, + 315, + 150 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "TANGOFLUX_VAE", + "link": 2 + }, + { + "name": "latents", + "type": "TANGOFLUX_LATENTS", + "link": 3 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "TangoFluxVAEDecodeAndPlay" + }, + "widgets_values": [ + 32, + "TangoFlux", + "wav", + true + ] + }, + { + "id": 1, + "type": "TangoFluxLoader", + "pos": [ + 100, + 230 + ], + "size": [ + 315, 102 ], "flags": {}, @@ -22,7 +59,7 @@ "name": "model", "type": "TANGOFLUX_MODEL", "links": [ - 11 + 1 ], "slot_index": 0 }, @@ -30,7 +67,7 @@ "name": "vae", "type": "TANGOFLUX_VAE", "links": [ - 15 + 2 ], "slot_index": 1 } @@ -44,47 +81,11 @@ ] }, { - "id": 13, - "type": "TangoFluxVAEDecodeAndPlay", - "pos": [ - 1060, - 320 - ], - "size": [ - 315, - 126 - ], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [ - { - "name": "vae", - "type": "TANGOFLUX_VAE", - "link": 15 - }, - { - "name": "latents", - "type": "TANGOFLUX_LATENTS", - "link": 14 - } - ], - "outputs": [], - "properties": { - "Node name for S&R": "TangoFluxVAEDecodeAndPlay" - }, - "widgets_values": [ - "TangoFlux", - "wav", - true - ] - }, - { - "id": 11, + "id": 2, "type": "TangoFluxSampler", "pos": [ - 620, - 320 + 450, + 230 ], "size": [ 400, @@ -97,7 +98,7 @@ { "name": "model", "type": "TANGOFLUX_MODEL", - "link": 11 + "link": 1 } ], "outputs": [ @@ -105,7 +106,7 @@ "name": "latents", "type": "TANGOFLUX_LATENTS", "links": [ - 14 + 3 ], "slot_index": 0 } @@ -118,7 +119,7 @@ 50, 3, 10, - 106139285587780, + 310686137950669, "randomize", 1 ] @@ -126,42 +127,39 @@ ], "links": [ [ - 11, - 10, + 1, + 1, 0, - 11, + 2, 0, "TANGOFLUX_MODEL" ], [ - 14, - 11, - 0, - 13, + 2, 1, - "TANGOFLUX_LATENTS" - ], - [ - 15, - 10, 1, - 13, + 3, 0, "TANGOFLUX_VAE" + ], + [ + 3, + 2, + 0, + 3, + 1, + "TANGOFLUX_LATENTS" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.9480295566502464, + "scale": 0.8630044843049327, "offset": [ - -200.83333333333337, - -102.2460379319304 + 95.83333333333337, + -8.663289166017137 ] - }, - "node_versions": { - "comfyui-tangoflux": "1.0.4" } }, "version": 0.4 diff --git a/nodes.py b/nodes.py index f4a8e17..bdd6e63 100644 --- a/nodes.py +++ b/nodes.py @@ -7,7 +7,6 @@ import torchaudio import re from diffusers import AutoencoderOobleck, FluxTransformer2DModel -from huggingface_hub import snapshot_download from comfy.utils import load_torch_file, ProgressBar import folder_paths @@ -66,16 +65,7 @@ class TangoFluxLoader: ): if self.model is None or self.enable_teacache != enable_teacache: - pbar = ProgressBar(6) - - snapshot_download( - repo_id="declare-lab/TangoFlux", - allow_patterns=["*.json", "*.safetensors"], - local_dir=tangoflux_path, - local_dir_use_symlinks=False, - ) - - pbar.update(1) + pbar = ProgressBar(4) log.info("Loading config") @@ -91,19 +81,12 @@ class TangoFluxLoader: ) text_encoder_path = os.path.join(text_encoder_path, text_encoder) - snapshot_download( - repo_id=config.get("text_encoder_name", "google/flan-t5-large"), - allow_patterns=["*.json", "*.safetensors", "*.model"], - local_dir=text_encoder_path, - local_dir_use_symlinks=False, - ) - - pbar.update(1) - log.info("Loading TangoFlux models") del self.model self.model = None + + torch.cuda.empty_cache() model_weights = load_torch_file( os.path.join(tangoflux_path, "tangoflux.safetensors"), @@ -116,7 +99,8 @@ class TangoFluxLoader: log.info("Enabling TeaCache") FluxTransformer2DModel.forward = teacache_forward else: - log.info("Disabling TeaCache") + if self.enable_teacache: + log.info("Disabling TeaCache") FluxTransformer2DModel.forward = self.original_forward model = TangoFlux(config, text_encoder_path) @@ -136,6 +120,9 @@ class TangoFluxLoader: self.model = model del model + + torch.cuda.empty_cache() + self.enable_teacache = enable_teacache self.rel_l1_thresh = rel_l1_thresh @@ -228,6 +215,7 @@ class TangoFluxVAEDecodeAndPlay: return { "required": { "vae": ("TANGOFLUX_VAE",), + "tile_size": ("INT", {"default": 32, "min": 8, "max": 128, "step": 8}), "latents": ("TANGOFLUX_LATENTS",), "filename_prefix": ("STRING", {"default": "TangoFlux"}), "format": ( @@ -245,21 +233,63 @@ class TangoFluxVAEDecodeAndPlay: FUNCTION = "play" DESCRIPTION = "Decoder and Player for TangoFlux" - def decode(self, vae, latents): + def decode_tiled(self, vae, latents, tile_size=32): + results = [] + + with torch.no_grad(): + for latent in latents: + torch.cuda.empty_cache() + + latent = latent.unsqueeze(0).transpose(2, 1) + + decoded_tiles = [] + for i in range(0, latent.size(2), tile_size): + tile = latent[:, :, i:i + tile_size] + + decoded_tile = vae.decode(tile).sample.cpu() + decoded_tiles.append(decoded_tile) + + del decoded_tile + torch.cuda.empty_cache() + + decoded_latent = torch.cat(decoded_tiles, dim=2) + results.append(decoded_latent) + + results = torch.cat(results, dim=0) + return results + + + def decode(self, vae, latents, tile_size=32): results = [] - for latent in latents: - decoded = vae.decode(latent.unsqueeze(0).transpose(2, 1)).sample.cpu() - results.append(decoded) + try: + with torch.no_grad(): + try: + print("Inducing OOM for testing...") + oom_tensor = torch.empty((1610612736,), dtype=torch.float32, device="cuda") + except torch.OutOfMemoryError: + print("OOM successfully induced for testing.") + raise torch.OutOfMemoryError + + for latent in latents: + torch.cuda.empty_cache() - results = torch.cat(results, dim=0) + decoded = vae.decode(latent.unsqueeze(0).transpose(2, 1)).sample.cpu() + results.append(decoded) - return results + results = torch.cat(results, dim=0) + return results + + except torch.OutOfMemoryError: + torch.cuda.empty_cache() + log.warning("OOM encountered. Falling back to tiled decoding.") + return self.decode_tiled(vae, latents, tile_size) def play( self, vae, latents, + tile_size=32, filename_prefix="TangoFlux", format="wav", save_output=True, @@ -285,7 +315,7 @@ class TangoFluxVAEDecodeAndPlay: ) os.makedirs(full_output_folder, exist_ok=True) - + pbar.update(1) duration = latents["duration"] @@ -295,8 +325,8 @@ class TangoFluxVAEDecodeAndPlay: log.info("Decoding Tangoflux latents") - waves = self.decode(vae, latents) - + waves = self.decode(vae, latents, tile_size) + pbar.update(1) for wave in waves: @@ -312,7 +342,7 @@ class TangoFluxVAEDecodeAndPlay: counter += 1 audios.append({"filename": file, "subfolder": subfolder, "type": type}) - + pbar.update(1) return {