From c50413ff1b63c6a1293e75a23e55810142801f43 Mon Sep 17 00:00:00 2001 From: LucipherDev Date: Sat, 11 Jan 2025 00:11:07 +0530 Subject: [PATCH] Bug Fixes --- example_workflow.json | 2 +- install.py | 44 +++++++++++++-------------- nodes.py | 70 +++++++++++++++++++++---------------------- pyproject.toml | 2 +- server.py | 5 ++-- tangoflux/model.py | 8 ++--- 6 files changed, 65 insertions(+), 66 deletions(-) diff --git a/example_workflow.json b/example_workflow.json index b784106..aceed95 100644 --- a/example_workflow.json +++ b/example_workflow.json @@ -161,7 +161,7 @@ ] }, "node_versions": { - "comfyui-tangoflux": "1.0.2" + "comfyui-tangoflux": "1.0.4" } }, "version": 0.4 diff --git a/install.py b/install.py index 1e59f18..906253e 100644 --- a/install.py +++ b/install.py @@ -25,26 +25,27 @@ except: download_models = False try: - log.info(f"Installing requirements") + log.info("Installing requirements") subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", f"{EXT_PATH}/requirements.txt", "--no-warn-script-location"]) if download_models: from huggingface_hub import snapshot_download + + log.info("Downloading Necessary models") try: - if not os.path.exists(TANGOFLUX_DIR): - log.info(f"Downloading TangoFlux models to: {TANGOFLUX_DIR}") - snapshot_download( - repo_id="declare-lab/TangoFlux", - allow_patterns=["*.json", "*.safetensors"], - local_dir=TANGOFLUX_DIR, - local_dir_use_symlinks=False, - ) + log.info(f"Downloading TangoFlux models to: {TANGOFLUX_DIR}") + snapshot_download( + repo_id="declare-lab/TangoFlux", + allow_patterns=["*.json", "*.safetensors"], + local_dir=TANGOFLUX_DIR, + local_dir_use_symlinks=False, + ) except Exception: traceback.print_exc() - log.error(f"Failed to download TangoFlux models") + log.error("Failed to download TangoFlux models") - log.info(f"Loading config") + log.info("Loading config") with open(os.path.join(TANGOFLUX_DIR, "config.json"), "r") as f: config = json.load(f) @@ -53,20 +54,19 @@ try: text_encoder = re.sub(r'[<>:"/\\|?*]', '-', config.get("text_encoder_name", "google/flan-t5-large")) text_encoder_path = os.path.join(TEXT_ENCODER_DIR, text_encoder) - if not os.path.exists(text_encoder_path): - log.info(f"Downloading text encoders to: {text_encoder_path}") - 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, - ) + log.info(f"Downloading text encoders to: {text_encoder_path}") + 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, + ) except Exception: traceback.print_exc() - log.error(f"Failed to download text encoders") + log.error("Failed to download text encoders") - log.info(f"TangoFlux Installation completed") + log.info("TangoFlux Installation completed") except Exception: traceback.print_exc() - log.error(f"TangoFlux Installation failed") \ No newline at end of file + log.error("TangoFlux Installation failed") \ No newline at end of file diff --git a/nodes.py b/nodes.py index 38807fe..f4a8e17 100644 --- a/nodes.py +++ b/nodes.py @@ -48,7 +48,7 @@ class TangoFluxLoader: CATEGORY = "TangoFlux" FUNCTION = "load_tangoflux" DESCRIPTION = "Load TangoFlux model" - + def __init__(self): self.model = None self.vae = None @@ -65,18 +65,15 @@ class TangoFluxLoader: device="cuda", ): if self.model is None or self.enable_teacache != enable_teacache: - del self.model - + pbar = ProgressBar(6) - if not os.path.exists(tangoflux_path): - log.info(f"Downloading TangoFlux models to: {tangoflux_path}") - snapshot_download( - repo_id="declare-lab/TangoFlux", - allow_patterns=["*.json", "*.safetensors"], - local_dir=tangoflux_path, - local_dir_use_symlinks=False, - ) + snapshot_download( + repo_id="declare-lab/TangoFlux", + allow_patterns=["*.json", "*.safetensors"], + local_dir=tangoflux_path, + local_dir_use_symlinks=False, + ) pbar.update(1) @@ -94,18 +91,19 @@ class TangoFluxLoader: ) text_encoder_path = os.path.join(text_encoder_path, text_encoder) - if not os.path.exists(text_encoder_path): - log.info(f"Downloading text encoders to: {text_encoder_path}") - 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, - ) + 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 model_weights = load_torch_file( os.path.join(tangoflux_path, "tangoflux.safetensors"), @@ -113,7 +111,7 @@ class TangoFluxLoader: ) pbar.update(1) - + if enable_teacache: log.info("Enabling TeaCache") FluxTransformer2DModel.forward = teacache_forward @@ -121,24 +119,26 @@ class TangoFluxLoader: log.info("Disabling TeaCache") FluxTransformer2DModel.forward = self.original_forward - self.model = TangoFlux(config, text_encoder_path) + model = TangoFlux(config, text_encoder_path) - self.model.load_state_dict(model_weights, strict=False) - self.model.to(device) + model.load_state_dict(model_weights, strict=False) + model.to(device) if enable_teacache: - self.model.transformer.__class__.enable_teacache = True - self.model.transformer.__class__.cnt = 0 - self.model.transformer.__class__.rel_l1_thresh = rel_l1_thresh - self.model.transformer.__class__.accumulated_rel_l1_distance = 0 - self.model.transformer.__class__.previous_modulated_input = None - self.model.transformer.__class__.previous_residual = None + model.transformer.__class__.enable_teacache = True + model.transformer.__class__.cnt = 0 + model.transformer.__class__.rel_l1_thresh = rel_l1_thresh + model.transformer.__class__.accumulated_rel_l1_distance = 0 + model.transformer.__class__.previous_modulated_input = None + model.transformer.__class__.previous_residual = None pbar.update(1) - + + self.model = model + del model self.enable_teacache = enable_teacache self.rel_l1_thresh = rel_l1_thresh - + if self.vae is None: log.info("Loading TangoFlux VAE") @@ -148,14 +148,14 @@ class TangoFluxLoader: self.vae = AutoencoderOobleck() self.vae.load_state_dict(vae_weights) self.vae.to(device) - + pbar.update(1) - + if self.enable_teacache == True and self.rel_l1_thresh != rel_l1_thresh: self.model.transformer.__class__.rel_l1_thresh = rel_l1_thresh - + self.rel_l1_thresh = rel_l1_thresh - + return (self.model, self.vae) diff --git a/pyproject.toml b/pyproject.toml index d91a68c..13abdd3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-tangoflux" description = "ComfyUI Custom Nodes for 'TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching'. This generates high-quality 44.1kHz audio up to 30 seconds using just a text prompt." -version = "1.0.3" +version = "1.0.4" license = {file = "LICENSE"} dependencies = ["torchaudio", "diffusers", "accelerate", "tqdm"] diff --git a/server.py b/server.py index cff2a7d..a8dcd96 100644 --- a/server.py +++ b/server.py @@ -4,10 +4,11 @@ import folder_paths web = server.web + @server.PromptServer.instance.routes.get("/tangoflux/playaudio") async def play_audio(request): query = request.rel_url.query - + filename = query.get("filename", None) if filename is None: @@ -43,7 +44,7 @@ async def play_audio(request): } content_type = content_types.get(ext, None) - + if content_type is None: return web.Response(status=400) diff --git a/tangoflux/model.py b/tangoflux/model.py index b44cec8..9604fa8 100644 --- a/tangoflux/model.py +++ b/tangoflux/model.py @@ -3,9 +3,9 @@ import torch from diffusers import FluxTransformer2DModel from torch import nn from diffusers import FlowMatchEulerDiscreteScheduler -import copy import numpy as np from tqdm import tqdm +import warnings from typing import Optional, Union, List, Any, Dict from math import pi @@ -73,7 +73,7 @@ def teacache_forward( joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None ): - logger.warning( + warnings.warn( "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective." ) hidden_states = self.x_embedder(hidden_states) @@ -420,11 +420,9 @@ class TangoFlux(nn.Module): self.num_attention_heads = config.get("num_attention_heads", 8) self.audio_seq_len = config.get("audio_seq_len", 645) self.max_duration = config.get("max_duration", 30) - self.uncondition = config.get("uncondition", False) self.text_encoder_name = config.get("text_encoder_name", "google/flan-t5-large") self.noise_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) - self.noise_scheduler_copy = copy.deepcopy(self.noise_scheduler) self.max_text_seq_len = 64 self.text_encoder = T5EncoderModel.from_pretrained( text_encoder_dir if text_encoder_dir is not None else self.text_encoder_name @@ -621,7 +619,7 @@ class TangoFlux(nn.Module): noise_pred = self.transformer( hidden_states=latents_input, - # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing) + # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transformer model (we should not keep it but I want to keep the inputs same for the model for testing) timestep=torch.tensor([t / 1000], device=device), guidance=None, pooled_projections=pooled_projection,