Bug Fixes

2025-01-11 00:11:07 +05:30
parent 8469da6a7d
commit c50413ff1b
6 changed files with 65 additions and 66 deletions
@@ -161,7 +161,7 @@
      ]
    },
    "node_versions": {
-      "comfyui-tangoflux": "1.0.2"
+      "comfyui-tangoflux": "1.0.4"
    }
  },
  "version": 0.4
@@ -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")
+    log.error("TangoFlux Installation failed")
@@ -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)


@@ -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"]

@@ -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)

@@ -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,