fix: only add timestamps to browser-previewed outputs

Reverted timestamp addition for non-previewed files: - SaveLatent (.latent) - not previewed in browser - CheckpointSave, CLIPSave, VAESave (.safetensors) - model files - ExtractAndSaveLoRA, SaveLoRA (.safetensors) - LoRA files Kept timestamps for browser-previewed content: - Images (PNG, SVG) - Videos (WebM, MP4) - Audio - 3D models (GLB) Amp-Thread-ID: https://ampcode.com/threads/T-019c1be0-7238-71ec-9c0b-2d4468d61202
fix: convert tests to unittest, remove unused import
2026-02-01 18:03:41 -08:00 · 2026-01-31 22:41:33 -08:00 · 2026-01-31 22:30:57 -08:00
221 changed files with 1475 additions and 12094 deletions
@@ -1,114 +0,0 @@
-# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json
-language: "en-US"
-early_access: false
-
-reviews:
-  profile: "chill"
-  request_changes_workflow: false
-  high_level_summary: false
-  poem: false
-  review_status: false
-  review_details: false
-  commit_status: true
-  collapse_walkthrough: true
-  changed_files_summary: false
-  sequence_diagrams: false
-  estimate_code_review_effort: false
-  assess_linked_issues: false
-  related_issues: false
-  related_prs: false
-  suggested_labels: false
-  auto_apply_labels: false
-  suggested_reviewers: false
-  auto_assign_reviewers: false
-  in_progress_fortune: false
-  enable_prompt_for_ai_agents: true
-
-  path_filters:
-    - "!comfy_api_nodes/apis/**"
-    - "!**/generated/*.pyi"
-    - "!.ci/**"
-    - "!script_examples/**"
-    - "!**/__pycache__/**"
-    - "!**/*.ipynb"
-    - "!**/*.png"
-    - "!**/*.bat"
-
-  path_instructions:
-    - path: "comfy/**"
-      instructions: |
-        Core ML/diffusion engine. Focus on:
-        - Backward compatibility (breaking changes affect all custom nodes)
-        - Memory management and GPU resource handling
-        - Performance implications in hot paths
-        - Thread safety for concurrent execution
-    - path: "comfy_api_nodes/**"
-      instructions: |
-        Third-party API integration nodes. Focus on:
-        - No hardcoded API keys or secrets
-        - Proper error handling for API failures (timeouts, rate limits, auth errors)
-        - Correct Pydantic model usage
-        - Security of user data passed to external APIs
-    - path: "comfy_extras/**"
-      instructions: |
-        Community-contributed extra nodes. Focus on:
-        - Consistency with node patterns (INPUT_TYPES, RETURN_TYPES, FUNCTION, CATEGORY)
-        - No breaking changes to existing node interfaces
-    - path: "comfy_execution/**"
-      instructions: |
-        Execution engine (graph execution, caching, jobs). Focus on:
-        - Caching correctness
-        - Concurrent execution safety
-        - Graph validation edge cases
-    - path: "nodes.py"
-      instructions: |
-        Core node definitions (2500+ lines). Focus on:
-        - Backward compatibility of NODE_CLASS_MAPPINGS
-        - Consistency of INPUT_TYPES return format
-    - path: "alembic_db/**"
-      instructions: |
-        Database migrations. Focus on:
-        - Migration safety and rollback support
-        - Data preservation during schema changes
-
-  auto_review:
-    enabled: true
-    auto_incremental_review: true
-    drafts: true
-
-  finishing_touches:
-    docstrings:
-      enabled: false
-    unit_tests:
-      enabled: false
-
-  tools:
-    ruff:
-      enabled: true
-    pylint:
-      enabled: false
-    flake8:
-      enabled: false
-    gitleaks:
-      enabled: true
-    shellcheck:
-      enabled: false
-    markdownlint:
-      enabled: false
-    yamllint:
-      enabled: false
-    languagetool:
-      enabled: false
-    github-checks:
-      enabled: true
-      timeout_ms: 90000
-    ast-grep:
-      essential_rules: true
-
-chat:
-  auto_reply: true
-
-knowledge_base:
-  opt_out: false
-  learnings:
-    scope: "auto"
@@ -7,8 +7,6 @@ on:
 jobs:
  send-webhook:
    runs-on: ubuntu-latest
-    env:
-      DESKTOP_REPO_DISPATCH_TOKEN: ${{ secrets.DESKTOP_REPO_DISPATCH_TOKEN }}
    steps:
      - name: Send release webhook
        env:
@@ -108,37 +106,3 @@ jobs:
            --fail --silent --show-error
          
          echo "✅ Release webhook sent successfully"
-
-      - name: Send repository dispatch to desktop
-        env:
-          DISPATCH_TOKEN: ${{ env.DESKTOP_REPO_DISPATCH_TOKEN }}
-          RELEASE_TAG: ${{ github.event.release.tag_name }}
-          RELEASE_URL: ${{ github.event.release.html_url }}
-        run: |
-          set -euo pipefail
-
-          if [ -z "${DISPATCH_TOKEN:-}" ]; then
-            echo "::error::DESKTOP_REPO_DISPATCH_TOKEN is required but not set."
-            exit 1
-          fi
-
-          PAYLOAD="$(jq -n \
-            --arg release_tag "$RELEASE_TAG" \
-            --arg release_url "$RELEASE_URL" \
-            '{
-              event_type: "comfyui_release_published",
-              client_payload: {
-                release_tag: $release_tag,
-                release_url: $release_url
-              }
-            }')"
-
-          curl -fsSL \
-            -X POST \
-            -H "Accept: application/vnd.github+json" \
-            -H "Content-Type: application/json" \
-            -H "Authorization: Bearer ${DISPATCH_TOKEN}" \
-            https://api.github.com/repos/Comfy-Org/desktop/dispatches \
-            -d "$PAYLOAD"
-
-          echo "✅ Dispatched ComfyUI release ${RELEASE_TAG} to Comfy-Org/desktop"
@@ -11,7 +11,7 @@ extra_model_paths.yaml
 /.vs
 .vscode/
 .idea/
-venv*/
+venv/
 .venv/
 /web/extensions/*
 !/web/extensions/logging.js.example
@@ -227,11 +227,11 @@ Put your VAE in: models/vae

 AMD users can install rocm and pytorch with pip if you don't have it already installed, this is the command to install the stable version:

-```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.1```
+```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.4```

-This is the command to install the nightly with ROCm 7.2 which might have some performance improvements:
+This is the command to install the nightly with ROCm 7.1 which might have some performance improvements:

-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.2```
+```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.1```


 ### AMD GPUs (Experimental: Windows and Linux), RDNA 3, 3.5 and 4 only.
@@ -1,105 +0,0 @@
-from __future__ import annotations
-
-from aiohttp import web
-
-from typing import TYPE_CHECKING, TypedDict
-if TYPE_CHECKING:
-    from comfy_api.latest._io_public import NodeReplace
-
-from comfy_execution.graph_utils import is_link
-import nodes
-
-class NodeStruct(TypedDict):
-    inputs: dict[str, str | int | float | bool | tuple[str, int]]
-    class_type: str
-    _meta: dict[str, str]
-
-def copy_node_struct(node_struct: NodeStruct, empty_inputs: bool = False) -> NodeStruct:
-    new_node_struct = node_struct.copy()
-    if empty_inputs:
-        new_node_struct["inputs"] = {}
-    else:
-        new_node_struct["inputs"] = node_struct["inputs"].copy()
-    new_node_struct["_meta"] = node_struct["_meta"].copy()
-    return new_node_struct
-
-
-class NodeReplaceManager:
-    """Manages node replacement registrations."""
-
-    def __init__(self):
-        self._replacements: dict[str, list[NodeReplace]] = {}
-
-    def register(self, node_replace: NodeReplace):
-        """Register a node replacement mapping."""
-        self._replacements.setdefault(node_replace.old_node_id, []).append(node_replace)
-
-    def get_replacement(self, old_node_id: str) -> list[NodeReplace] | None:
-        """Get replacements for an old node ID."""
-        return self._replacements.get(old_node_id)
-
-    def has_replacement(self, old_node_id: str) -> bool:
-        """Check if a replacement exists for an old node ID."""
-        return old_node_id in self._replacements
-
-    def apply_replacements(self, prompt: dict[str, NodeStruct]):
-        connections: dict[str, list[tuple[str, str, int]]] = {}
-        need_replacement: set[str] = set()
-        for node_number, node_struct in prompt.items():
-            class_type = node_struct["class_type"]
-            # need replacement if not in NODE_CLASS_MAPPINGS and has replacement
-            if class_type not in nodes.NODE_CLASS_MAPPINGS.keys() and self.has_replacement(class_type):
-                need_replacement.add(node_number)
-            # keep track of connections
-            for input_id, input_value in node_struct["inputs"].items():
-                if is_link(input_value):
-                    conn_number = input_value[0]
-                    connections.setdefault(conn_number, []).append((node_number, input_id, input_value[1]))
-        for node_number in need_replacement:
-            node_struct = prompt[node_number]
-            class_type = node_struct["class_type"]
-            replacements = self.get_replacement(class_type)
-            if replacements is None:
-                continue
-            # just use the first replacement
-            replacement = replacements[0]
-            new_node_id = replacement.new_node_id
-            # if replacement is not a valid node, skip trying to replace it as will only cause confusion
-            if new_node_id not in nodes.NODE_CLASS_MAPPINGS.keys():
-                continue
-            # first, replace node id (class_type)
-            new_node_struct = copy_node_struct(node_struct, empty_inputs=True)
-            new_node_struct["class_type"] = new_node_id
-            # TODO: consider replacing display_name in _meta as well for error reporting purposes; would need to query node schema
-            # second, replace inputs
-            if replacement.input_mapping is not None:
-                for input_map in replacement.input_mapping:
-                    if "set_value" in input_map:
-                        new_node_struct["inputs"][input_map["new_id"]] = input_map["set_value"]
-                    elif "old_id" in input_map:
-                        new_node_struct["inputs"][input_map["new_id"]] = node_struct["inputs"][input_map["old_id"]]
-            # finalize input replacement
-            prompt[node_number] = new_node_struct
-            # third, replace outputs
-            if replacement.output_mapping is not None:
-                # re-mapping outputs requires changing the input values of nodes that receive connections from this one
-                if node_number in connections:
-                    for conns in connections[node_number]:
-                        conn_node_number, conn_input_id, old_output_idx = conns
-                        for output_map in replacement.output_mapping:
-                            if output_map["old_idx"] == old_output_idx:
-                                new_output_idx = output_map["new_idx"]
-                                previous_input = prompt[conn_node_number]["inputs"][conn_input_id]
-                                previous_input[1] = new_output_idx
-
-    def as_dict(self):
-        """Serialize all replacements to dict."""
-        return {
-            k: [v.as_dict() for v in v_list]
-            for k, v_list in self._replacements.items()
-        }
-
-    def add_routes(self, routes):
-        @routes.get("/node_replacements")
-        async def get_node_replacements(request):
-            return web.json_response(self.as_dict())
@@ -1,44 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform float u_float0; // Brightness slider -100..100
-uniform float u_float1; // Contrast slider -100..100
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-const float MID_GRAY = 0.18;  // 18% reflectance
-
-// sRGB gamma 2.2 approximation
-vec3 srgbToLinear(vec3 c) {
-    return pow(max(c, 0.0), vec3(2.2));
-}
-
-vec3 linearToSrgb(vec3 c) {
-    return pow(max(c, 0.0), vec3(1.0/2.2));
-}
-
-float mapBrightness(float b) {
-    return clamp(b / 100.0, -1.0, 1.0);
-}
-
-float mapContrast(float c) {
-    return clamp(c / 100.0 + 1.0, 0.0, 2.0);
-}
-
-void main() {
-    vec4 orig = texture(u_image0, v_texCoord);
-
-    float brightness = mapBrightness(u_float0);
-    float contrast   = mapContrast(u_float1);
-
-    vec3 lin = srgbToLinear(orig.rgb);
-
-    lin = (lin - MID_GRAY) * contrast + brightness + MID_GRAY;
-
-    // Convert back to sRGB
-    vec3 result = linearToSrgb(clamp(lin, 0.0, 1.0));
-
-    fragColor = vec4(result, orig.a);
-}
@@ -1,72 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform vec2 u_resolution;
-uniform int u_int0;      // Mode
-uniform float u_float0;  // Amount (0 to 100)
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-const int MODE_LINEAR   = 0;
-const int MODE_RADIAL   = 1;
-const int MODE_BARREL   = 2;
-const int MODE_SWIRL    = 3;
-const int MODE_DIAGONAL = 4;
-
-const float AMOUNT_SCALE = 0.0005;
-const float RADIAL_MULT = 4.0;
-const float BARREL_MULT = 8.0;
-const float INV_SQRT2 = 0.70710678118;
-
-void main() {
-    vec2 uv = v_texCoord;
-    vec4 original = texture(u_image0, uv);
-
-    float amount = u_float0 * AMOUNT_SCALE;
-
-    if (amount < 0.000001) {
-        fragColor = original;
-        return;
-    }
-
-    // Aspect-corrected coordinates for circular effects
-    float aspect = u_resolution.x / u_resolution.y;
-    vec2 centered = uv - 0.5;
-    vec2 corrected = vec2(centered.x * aspect, centered.y);
-    float r = length(corrected);
-    vec2 dir = r > 0.0001 ? corrected / r : vec2(0.0);
-    vec2 offset = vec2(0.0);
-
-    if (u_int0 == MODE_LINEAR) {
-        // Horizontal shift (no aspect correction needed)
-        offset = vec2(amount, 0.0);
-    }
-    else if (u_int0 == MODE_RADIAL) {
-        // Outward from center, stronger at edges
-        offset = dir * r * amount * RADIAL_MULT;
-        offset.x /= aspect;  // Convert back to UV space
-    }
-    else if (u_int0 == MODE_BARREL) {
-        // Lens distortion simulation (r² falloff)
-        offset = dir * r * r * amount * BARREL_MULT;
-        offset.x /= aspect;  // Convert back to UV space
-    }
-    else if (u_int0 == MODE_SWIRL) {
-        // Perpendicular to radial (rotational aberration)
-        vec2 perp = vec2(-dir.y, dir.x);
-        offset = perp * r * amount * RADIAL_MULT;
-        offset.x /= aspect;  // Convert back to UV space
-    }
-    else if (u_int0 == MODE_DIAGONAL) {
-        // 45° offset (no aspect correction needed)
-        offset = vec2(amount, amount) * INV_SQRT2;
-    }
-    
-    float red = texture(u_image0, uv + offset).r;
-    float green = original.g;
-    float blue = texture(u_image0, uv - offset).b;
-    
-    fragColor = vec4(red, green, blue, original.a);
-}
@@ -1,78 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform float u_float0; // temperature (-100 to 100)
-uniform float u_float1; // tint (-100 to 100)
-uniform float u_float2; // vibrance (-100 to 100)
-uniform float u_float3; // saturation (-100 to 100)
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-const float INPUT_SCALE = 0.01;
-const float TEMP_TINT_PRIMARY = 0.3;
-const float TEMP_TINT_SECONDARY = 0.15;
-const float VIBRANCE_BOOST = 2.0;
-const float SATURATION_BOOST = 2.0;
-const float SKIN_PROTECTION = 0.5;
-const float EPSILON = 0.001;
-const vec3 LUMA_WEIGHTS = vec3(0.299, 0.587, 0.114);
-
-void main() {
-    vec4 tex = texture(u_image0, v_texCoord);
-    vec3 color = tex.rgb;
-    
-    // Scale inputs: -100/100 → -1/1
-    float temperature = u_float0 * INPUT_SCALE;
-    float tint = u_float1 * INPUT_SCALE;
-    float vibrance = u_float2 * INPUT_SCALE;
-    float saturation = u_float3 * INPUT_SCALE;
-    
-    // Temperature (warm/cool): positive = warm, negative = cool
-    color.r += temperature * TEMP_TINT_PRIMARY;
-    color.b -= temperature * TEMP_TINT_PRIMARY;
-    
-    // Tint (green/magenta): positive = green, negative = magenta
-    color.g += tint * TEMP_TINT_PRIMARY;
-    color.r -= tint * TEMP_TINT_SECONDARY;
-    color.b -= tint * TEMP_TINT_SECONDARY;
-    
-    // Single clamp after temperature/tint
-    color = clamp(color, 0.0, 1.0);
-    
-    // Vibrance with skin protection
-    if (vibrance != 0.0) {
-        float maxC = max(color.r, max(color.g, color.b));
-        float minC = min(color.r, min(color.g, color.b));
-        float sat = maxC - minC;
-        float gray = dot(color, LUMA_WEIGHTS);
-        
-        if (vibrance < 0.0) {
-            // Desaturate: -100 → gray
-            color = mix(vec3(gray), color, 1.0 + vibrance);
-        } else {
-            // Boost less saturated colors more
-            float vibranceAmt = vibrance * (1.0 - sat);
-            
-            // Branchless skin tone protection
-            float isWarmTone = step(color.b, color.g) * step(color.g, color.r);
-            float warmth = (color.r - color.b) / max(maxC, EPSILON);
-            float skinTone = isWarmTone * warmth * sat * (1.0 - sat);
-            vibranceAmt *= (1.0 - skinTone * SKIN_PROTECTION);
-            
-            color = mix(vec3(gray), color, 1.0 + vibranceAmt * VIBRANCE_BOOST);
-        }
-    }
-    
-    // Saturation
-    if (saturation != 0.0) {
-        float gray = dot(color, LUMA_WEIGHTS);
-        float satMix = saturation < 0.0
-            ? 1.0 + saturation                      // -100 → gray
-            : 1.0 + saturation * SATURATION_BOOST;  // +100 → 3x boost
-        color = mix(vec3(gray), color, satMix);
-    }
-    
-    fragColor = vec4(clamp(color, 0.0, 1.0), tex.a);
-}
@@ -1,94 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform float u_float0;   // Blur radius (0–20, default ~5)
-uniform float u_float1;   // Edge threshold (0–100, default ~30)
-uniform int u_int0;       // Step size (0/1 = every pixel, 2+ = skip pixels)
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-const int MAX_RADIUS = 20;
-const float EPSILON = 0.0001;
-
-// Perceptual luminance
-float getLuminance(vec3 rgb) {
-    return dot(rgb, vec3(0.299, 0.587, 0.114));
-}
-
-vec4 bilateralFilter(vec2 uv, vec2 texelSize, int radius,
-                     float sigmaSpatial, float sigmaColor)
-{
-    vec4 center = texture(u_image0, uv);
-    vec3 centerRGB = center.rgb;
-
-    float invSpatial2 = -0.5 / (sigmaSpatial * sigmaSpatial);
-    float invColor2   = -0.5 / (sigmaColor * sigmaColor + EPSILON);
-
-    vec3 sumRGB = vec3(0.0);
-    float sumWeight = 0.0;
-
-    int step = max(u_int0, 1);
-    float radius2 = float(radius * radius);
-
-    for (int dy = -MAX_RADIUS; dy <= MAX_RADIUS; dy++) {
-        if (dy < -radius || dy > radius) continue;
-        if (abs(dy) % step != 0) continue;
-
-        for (int dx = -MAX_RADIUS; dx <= MAX_RADIUS; dx++) {
-            if (dx < -radius || dx > radius) continue;
-            if (abs(dx) % step != 0) continue;
-
-            vec2 offset = vec2(float(dx), float(dy));
-            float dist2 = dot(offset, offset);
-            if (dist2 > radius2) continue;
-
-            vec3 sampleRGB = texture(u_image0, uv + offset * texelSize).rgb;
-
-            // Spatial Gaussian
-            float spatialWeight = exp(dist2 * invSpatial2);
-
-            // Perceptual color distance (weighted RGB)
-            vec3 diff = sampleRGB - centerRGB;
-            float colorDist = dot(diff * diff, vec3(0.299, 0.587, 0.114));
-            float colorWeight = exp(colorDist * invColor2);
-
-            float w = spatialWeight * colorWeight;
-            sumRGB += sampleRGB * w;
-            sumWeight += w;
-        }
-    }
-
-    vec3 resultRGB = sumRGB / max(sumWeight, EPSILON);
-    return vec4(resultRGB, center.a); // preserve center alpha
-}
-
-void main() {
-    vec2 texelSize = 1.0 / vec2(textureSize(u_image0, 0));
-
-    float radiusF = clamp(u_float0, 0.0, float(MAX_RADIUS));
-    int radius = int(radiusF + 0.5);
-
-    if (radius == 0) {
-        fragColor = texture(u_image0, v_texCoord);
-        return;
-    }
-
-    // Edge threshold → color sigma
-    // Squared curve for better low-end control
-    float t = clamp(u_float1, 0.0, 100.0) / 100.0;
-    t *= t;
-    float sigmaColor = mix(0.01, 0.5, t);
-
-    // Spatial sigma tied to radius
-    float sigmaSpatial = max(radiusF * 0.75, 0.5);
-
-    fragColor = bilateralFilter(
-        v_texCoord,
-        texelSize,
-        radius,
-        sigmaSpatial,
-        sigmaColor
-    );
-}
@@ -1,124 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform vec2 u_resolution;
-uniform float u_float0; // grain amount      [0.0 – 1.0]   typical: 0.2–0.8
-uniform float u_float1; // grain size        [0.3 – 3.0]   lower = finer grain
-uniform float u_float2; // color amount      [0.0 – 1.0]   0 = monochrome, 1 = RGB grain
-uniform float u_float3; // luminance bias    [0.0 – 1.0]   0 = uniform, 1 = shadows only
-uniform int   u_int0;   // noise mode        [0 or 1]      0 = smooth, 1 = grainy
-
-in vec2 v_texCoord;
-layout(location = 0) out vec4 fragColor0;
-
-// High-quality integer hash (pcg-like)
-uint pcg(uint v) {
-    uint state = v * 747796405u + 2891336453u;
-    uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
-    return (word >> 22u) ^ word;
-}
-
-// 2D -> 1D hash input
-uint hash2d(uvec2 p) {
-    return pcg(p.x + pcg(p.y));
-}
-
-// Hash to float [0, 1]
-float hashf(uvec2 p) {
-    return float(hash2d(p)) / float(0xffffffffu);
-}
-
-// Hash to float with offset (for RGB channels)
-float hashf(uvec2 p, uint offset) {
-    return float(pcg(hash2d(p) + offset)) / float(0xffffffffu);
-}
-
-// Convert uniform [0,1] to roughly Gaussian distribution
-// Using simple approximation: average of multiple samples
-float toGaussian(uvec2 p) {
-    float sum = hashf(p, 0u) + hashf(p, 1u) + hashf(p, 2u) + hashf(p, 3u);
-    return (sum - 2.0) * 0.7;  // Centered, scaled
-}
-
-float toGaussian(uvec2 p, uint offset) {
-    float sum = hashf(p, offset) + hashf(p, offset + 1u) 
-              + hashf(p, offset + 2u) + hashf(p, offset + 3u);
-    return (sum - 2.0) * 0.7;
-}
-
-// Smooth noise with better interpolation
-float smoothNoise(vec2 p) {
-    vec2 i = floor(p);
-    vec2 f = fract(p);
-    
-    // Quintic interpolation (less banding than cubic)
-    f = f * f * f * (f * (f * 6.0 - 15.0) + 10.0);
-    
-    uvec2 ui = uvec2(i);
-    float a = toGaussian(ui);
-    float b = toGaussian(ui + uvec2(1u, 0u));
-    float c = toGaussian(ui + uvec2(0u, 1u));
-    float d = toGaussian(ui + uvec2(1u, 1u));
-    
-    return mix(mix(a, b, f.x), mix(c, d, f.x), f.y);
-}
-
-float smoothNoise(vec2 p, uint offset) {
-    vec2 i = floor(p);
-    vec2 f = fract(p);
-    
-    f = f * f * f * (f * (f * 6.0 - 15.0) + 10.0);
-    
-    uvec2 ui = uvec2(i);
-    float a = toGaussian(ui, offset);
-    float b = toGaussian(ui + uvec2(1u, 0u), offset);
-    float c = toGaussian(ui + uvec2(0u, 1u), offset);
-    float d = toGaussian(ui + uvec2(1u, 1u), offset);
-    
-    return mix(mix(a, b, f.x), mix(c, d, f.x), f.y);
-}
-
-void main() {
-    vec4 color = texture(u_image0, v_texCoord);
-    
-    // Luminance (Rec.709)
-    float luma = dot(color.rgb, vec3(0.2126, 0.7152, 0.0722));
-    
-    // Grain UV (resolution-independent)
-    vec2 grainUV = v_texCoord * u_resolution / max(u_float1, 0.01);
-    uvec2 grainPixel = uvec2(grainUV);
-    
-    float g;
-    vec3 grainRGB;
-    
-    if (u_int0 == 1) {
-        // Grainy mode: pure hash noise (no interpolation = no banding)
-        g = toGaussian(grainPixel);
-        grainRGB = vec3(
-            toGaussian(grainPixel, 100u),
-            toGaussian(grainPixel, 200u),
-            toGaussian(grainPixel, 300u)
-        );
-    } else {
-        // Smooth mode: interpolated with quintic curve
-        g = smoothNoise(grainUV);
-        grainRGB = vec3(
-            smoothNoise(grainUV, 100u),
-            smoothNoise(grainUV, 200u),
-            smoothNoise(grainUV, 300u)
-        );
-    }
-    
-    // Luminance weighting (less grain in highlights)
-    float lumWeight = mix(1.0, 1.0 - luma, clamp(u_float3, 0.0, 1.0));
-    
-    // Strength
-    float strength = u_float0 * 0.15;
-    
-    // Color vs monochrome grain
-    vec3 grainColor = mix(vec3(g), grainRGB, clamp(u_float2, 0.0, 1.0));
-    
-    color.rgb += grainColor * strength * lumWeight;
-    fragColor0 = vec4(clamp(color.rgb, 0.0, 1.0), color.a);
-}
@@ -1,133 +0,0 @@
-#version 300 es
-precision mediump float;
-
-uniform sampler2D u_image0;
-uniform vec2 u_resolution;
-uniform int u_int0;      // Blend mode
-uniform int u_int1;      // Color tint
-uniform float u_float0;  // Intensity
-uniform float u_float1;  // Radius
-uniform float u_float2;  // Threshold
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-const int BLEND_ADD      = 0;
-const int BLEND_SCREEN   = 1;
-const int BLEND_SOFT     = 2;
-const int BLEND_OVERLAY  = 3;
-const int BLEND_LIGHTEN  = 4;
-
-const float GOLDEN_ANGLE = 2.39996323;
-const int MAX_SAMPLES = 48;
-const vec3 LUMA = vec3(0.299, 0.587, 0.114);
-
-float hash(vec2 p) {
-    p = fract(p * vec2(123.34, 456.21));
-    p += dot(p, p + 45.32);
-    return fract(p.x * p.y);
-}
-
-vec3 hexToRgb(int h) {
-    return vec3(
-        float((h >> 16) & 255),
-        float((h >> 8) & 255),
-        float(h & 255)
-    ) * (1.0 / 255.0);
-}
-
-vec3 blend(vec3 base, vec3 glow, int mode) {
-    if (mode == BLEND_SCREEN) {
-        return 1.0 - (1.0 - base) * (1.0 - glow);
-    }
-    if (mode == BLEND_SOFT) {
-        return mix(
-            base - (1.0 - 2.0 * glow) * base * (1.0 - base),
-            base + (2.0 * glow - 1.0) * (sqrt(base) - base),
-            step(0.5, glow)
-        );
-    }
-    if (mode == BLEND_OVERLAY) {
-        return mix(
-            2.0 * base * glow,
-            1.0 - 2.0 * (1.0 - base) * (1.0 - glow),
-            step(0.5, base)
-        );
-    }
-    if (mode == BLEND_LIGHTEN) {
-        return max(base, glow);
-    }
-    return base + glow;
-}
-
-void main() {
-    vec4 original = texture(u_image0, v_texCoord);
-    
-    float intensity = u_float0 * 0.05;
-    float radius = u_float1 * u_float1 * 0.012;
-    
-    if (intensity < 0.001 || radius < 0.1) {
-        fragColor = original;
-        return;
-    }
-    
-    float threshold = 1.0 - u_float2 * 0.01;
-    float t0 = threshold - 0.15;
-    float t1 = threshold + 0.15;
-    
-    vec2 texelSize = 1.0 / u_resolution;
-    float radius2 = radius * radius;
-    
-    float sampleScale = clamp(radius * 0.75, 0.35, 1.0);
-    int samples = int(float(MAX_SAMPLES) * sampleScale);
-    
-    float noise = hash(gl_FragCoord.xy);
-    float angleOffset = noise * GOLDEN_ANGLE;
-    float radiusJitter = 0.85 + noise * 0.3;
-    
-    float ca = cos(GOLDEN_ANGLE);
-    float sa = sin(GOLDEN_ANGLE);
-    vec2 dir = vec2(cos(angleOffset), sin(angleOffset));
-    
-    vec3 glow = vec3(0.0);
-    float totalWeight = 0.0;
-    
-    // Center tap
-    float centerMask = smoothstep(t0, t1, dot(original.rgb, LUMA));
-    glow += original.rgb * centerMask * 2.0;
-    totalWeight += 2.0;
-    
-    for (int i = 1; i < MAX_SAMPLES; i++) {
-        if (i >= samples) break;
-        
-        float fi = float(i);
-        float dist = sqrt(fi / float(samples)) * radius * radiusJitter;
-        
-        vec2 offset = dir * dist * texelSize;
-        vec3 c = texture(u_image0, v_texCoord + offset).rgb;
-        float mask = smoothstep(t0, t1, dot(c, LUMA));
-        
-        float w = 1.0 - (dist * dist) / (radius2 * 1.5);
-        w = max(w, 0.0);
-        w *= w;
-        
-        glow += c * mask * w;
-        totalWeight += w;
-        
-        dir = vec2(
-            dir.x * ca - dir.y * sa,
-            dir.x * sa + dir.y * ca
-        );
-    }
-    
-    glow *= intensity / max(totalWeight, 0.001);
-    
-    if (u_int1 > 0) {
-        glow *= hexToRgb(u_int1);
-    }
-    
-    vec3 result = blend(original.rgb, glow, u_int0);
-    result += (noise - 0.5) * (1.0 / 255.0);
-    
-    fragColor = vec4(clamp(result, 0.0, 1.0), original.a);
-}
@@ -1,222 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform int u_int0;      // Mode: 0=Master, 1=Reds, 2=Yellows, 3=Greens, 4=Cyans, 5=Blues, 6=Magentas, 7=Colorize
-uniform int u_int1;      // Color Space: 0=HSL, 1=HSB/HSV
-uniform float u_float0;  // Hue (-180 to 180)
-uniform float u_float1;  // Saturation (-100 to 100)
-uniform float u_float2;  // Lightness/Brightness (-100 to 100)
-uniform float u_float3;  // Overlap (0 to 100) - feathering between adjacent color ranges
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-// Color range modes
-const int MODE_MASTER   = 0;
-const int MODE_RED      = 1;
-const int MODE_YELLOW   = 2;
-const int MODE_GREEN    = 3;
-const int MODE_CYAN     = 4;
-const int MODE_BLUE     = 5;
-const int MODE_MAGENTA  = 6;
-const int MODE_COLORIZE = 7;
-
-// Color space modes
-const int COLORSPACE_HSL = 0;
-const int COLORSPACE_HSB = 1;
-
-const float EPSILON = 0.0001;
-
-//=============================================================================
-// RGB <-> HSL Conversions
-//=============================================================================
-
-vec3 rgb2hsl(vec3 c) {
-    float maxC = max(max(c.r, c.g), c.b);
-    float minC = min(min(c.r, c.g), c.b);
-    float delta = maxC - minC;
-
-    float h = 0.0;
-    float s = 0.0;
-    float l = (maxC + minC) * 0.5;
-
-    if (delta > EPSILON) {
-        s = l < 0.5
-            ? delta / (maxC + minC)
-            : delta / (2.0 - maxC - minC);
-
-        if (maxC == c.r) {
-            h = (c.g - c.b) / delta + (c.g < c.b ? 6.0 : 0.0);
-        } else if (maxC == c.g) {
-            h = (c.b - c.r) / delta + 2.0;
-        } else {
-            h = (c.r - c.g) / delta + 4.0;
-        }
-        h /= 6.0;
-    }
-
-    return vec3(h, s, l);
-}
-
-float hue2rgb(float p, float q, float t) {
-    t = fract(t);
-    if (t < 1.0/6.0) return p + (q - p) * 6.0 * t;
-    if (t < 0.5)       return q;
-    if (t < 2.0/3.0)   return p + (q - p) * (2.0/3.0 - t) * 6.0;
-    return p;
-}
-
-vec3 hsl2rgb(vec3 hsl) {
-    if (hsl.y < EPSILON) return vec3(hsl.z);
-
-    float q = hsl.z < 0.5
-        ? hsl.z * (1.0 + hsl.y)
-        : hsl.z + hsl.y - hsl.z * hsl.y;
-    float p = 2.0 * hsl.z - q;
-
-    return vec3(
-        hue2rgb(p, q, hsl.x + 1.0/3.0),
-        hue2rgb(p, q, hsl.x),
-        hue2rgb(p, q, hsl.x - 1.0/3.0)
-    );
-}
-
-vec3 rgb2hsb(vec3 c) {
-    float maxC = max(max(c.r, c.g), c.b);
-    float minC = min(min(c.r, c.g), c.b);
-    float delta = maxC - minC;
-
-    float h = 0.0;
-    float s = (maxC > EPSILON) ? delta / maxC : 0.0;
-    float b = maxC;
-
-    if (delta > EPSILON) {
-        if (maxC == c.r) {
-            h = (c.g - c.b) / delta + (c.g < c.b ? 6.0 : 0.0);
-        } else if (maxC == c.g) {
-            h = (c.b - c.r) / delta + 2.0;
-        } else {
-            h = (c.r - c.g) / delta + 4.0;
-        }
-        h /= 6.0;
-    }
-
-    return vec3(h, s, b);
-}
-
-vec3 hsb2rgb(vec3 hsb) {
-    vec3 rgb = clamp(abs(mod(hsb.x * 6.0 + vec3(0.0, 4.0, 2.0), 6.0) - 3.0) - 1.0, 0.0, 1.0);
-    return hsb.z * mix(vec3(1.0), rgb, hsb.y);
-}
-
-//=============================================================================
-// Color Range Weight Calculation
-//=============================================================================
-
-float hueDistance(float a, float b) {
-    float d = abs(a - b);
-    return min(d, 1.0 - d);
-}
-
-float getHueWeight(float hue, float center, float overlap) {
-    float baseWidth = 1.0 / 6.0;
-    float feather = baseWidth * overlap;
-
-    float d = hueDistance(hue, center);
-
-    float inner = baseWidth * 0.5;
-    float outer = inner + feather;
-
-    return 1.0 - smoothstep(inner, outer, d);
-}
-
-float getModeWeight(float hue, int mode, float overlap) {
-    if (mode == MODE_MASTER || mode == MODE_COLORIZE) return 1.0;
-
-    if (mode == MODE_RED) {
-        return max(
-            getHueWeight(hue, 0.0, overlap),
-            getHueWeight(hue, 1.0, overlap)
-        );
-    }
-
-    float center = float(mode - 1) / 6.0;
-    return getHueWeight(hue, center, overlap);
-}
-
-//=============================================================================
-// Adjustment Functions
-//=============================================================================
-
-float adjustLightness(float l, float amount) {
-    return amount > 0.0
-        ? l + (1.0 - l) * amount
-        : l + l * amount;
-}
-
-float adjustBrightness(float b, float amount) {
-    return clamp(b + amount, 0.0, 1.0);
-}
-
-float adjustSaturation(float s, float amount) {
-    return amount > 0.0
-        ? s + (1.0 - s) * amount
-        : s + s * amount;
-}
-
-vec3 colorize(vec3 rgb, float hue, float sat, float light) {
-    float lum = dot(rgb, vec3(0.299, 0.587, 0.114));
-    float l = adjustLightness(lum, light);
-
-    vec3 hsl = vec3(fract(hue), clamp(sat, 0.0, 1.0), clamp(l, 0.0, 1.0));
-    return hsl2rgb(hsl);
-}
-
-//=============================================================================
-// Main
-//=============================================================================
-
-void main() {
-    vec4 original = texture(u_image0, v_texCoord);
-
-    float hueShift   = u_float0 / 360.0;   // -180..180 -> -0.5..0.5
-    float satAmount  = u_float1 / 100.0;   // -100..100 -> -1..1
-    float lightAmount= u_float2 / 100.0;   // -100..100 -> -1..1
-    float overlap    = u_float3 / 100.0;   // 0..100 -> 0..1
-
-    vec3 result;
-
-    if (u_int0 == MODE_COLORIZE) {
-        result = colorize(original.rgb, hueShift, satAmount, lightAmount);
-        fragColor = vec4(result, original.a);
-        return;
-    }
-
-    vec3 hsx = (u_int1 == COLORSPACE_HSL)
-        ? rgb2hsl(original.rgb)
-        : rgb2hsb(original.rgb);
-
-    float weight = getModeWeight(hsx.x, u_int0, overlap);
-
-    if (u_int0 != MODE_MASTER && hsx.y < EPSILON) {
-        weight = 0.0;
-    }
-
-    if (weight > EPSILON) {
-        float h = fract(hsx.x + hueShift * weight);
-        float s = clamp(adjustSaturation(hsx.y, satAmount * weight), 0.0, 1.0);
-        float v = (u_int1 == COLORSPACE_HSL)
-            ? clamp(adjustLightness(hsx.z, lightAmount * weight), 0.0, 1.0)
-            : clamp(adjustBrightness(hsx.z, lightAmount * weight), 0.0, 1.0);
-
-        vec3 adjusted = vec3(h, s, v);
-        result = (u_int1 == COLORSPACE_HSL)
-            ? hsl2rgb(adjusted)
-            : hsb2rgb(adjusted);
-    } else {
-        result = original.rgb;
-    }
-
-    fragColor = vec4(result, original.a);
-}
@@ -1,111 +0,0 @@
-#version 300 es
-#pragma passes 2
-precision highp float;
-
-// Blur type constants
-const int BLUR_GAUSSIAN = 0;
-const int BLUR_BOX = 1;
-const int BLUR_RADIAL = 2;
-
-// Radial blur config
-const int RADIAL_SAMPLES = 12;
-const float RADIAL_STRENGTH = 0.0003;
-
-uniform sampler2D u_image0;
-uniform vec2 u_resolution;
-uniform int u_int0;      // Blur type (BLUR_GAUSSIAN, BLUR_BOX, BLUR_RADIAL)
-uniform float u_float0;  // Blur radius/amount
-uniform int u_pass;      // Pass index (0 = horizontal, 1 = vertical)
-
-in vec2 v_texCoord;
-layout(location = 0) out vec4 fragColor0;
-
-float gaussian(float x, float sigma) {
-    return exp(-(x * x) / (2.0 * sigma * sigma));
-}
-
-void main() {
-    vec2 texelSize = 1.0 / u_resolution;
-    float radius = max(u_float0, 0.0);
-
-    // Radial (angular) blur - single pass, doesn't use separable
-    if (u_int0 == BLUR_RADIAL) {
-        // Only execute on first pass
-        if (u_pass > 0) {
-            fragColor0 = texture(u_image0, v_texCoord);
-            return;
-        }
-
-        vec2 center = vec2(0.5);
-        vec2 dir = v_texCoord - center;
-        float dist = length(dir);
-
-        if (dist < 1e-4) {
-            fragColor0 = texture(u_image0, v_texCoord);
-            return;
-        }
-
-        vec4 sum = vec4(0.0);
-        float totalWeight = 0.0;
-        float angleStep = radius * RADIAL_STRENGTH;
-
-        dir /= dist;
-
-        float cosStep = cos(angleStep);
-        float sinStep = sin(angleStep);
-
-        float negAngle = -float(RADIAL_SAMPLES) * angleStep;
-        vec2 rotDir = vec2(
-            dir.x * cos(negAngle) - dir.y * sin(negAngle),
-            dir.x * sin(negAngle) + dir.y * cos(negAngle)
-        );
-
-        for (int i = -RADIAL_SAMPLES; i <= RADIAL_SAMPLES; i++) {
-            vec2 uv = center + rotDir * dist;
-            float w = 1.0 - abs(float(i)) / float(RADIAL_SAMPLES);
-            sum += texture(u_image0, uv) * w;
-            totalWeight += w;
-
-            rotDir = vec2(
-                rotDir.x * cosStep - rotDir.y * sinStep,
-                rotDir.x * sinStep + rotDir.y * cosStep
-            );
-        }
-
-        fragColor0 = sum / max(totalWeight, 0.001);
-        return;
-    }
-
-    // Separable Gaussian / Box blur
-    int samples = int(ceil(radius));
-
-    if (samples == 0) {
-        fragColor0 = texture(u_image0, v_texCoord);
-        return;
-    }
-
-    // Direction: pass 0 = horizontal, pass 1 = vertical
-    vec2 dir = (u_pass == 0) ? vec2(1.0, 0.0) : vec2(0.0, 1.0);
-
-    vec4 color = vec4(0.0);
-    float totalWeight = 0.0;
-    float sigma = radius / 2.0;
-
-    for (int i = -samples; i <= samples; i++) {
-        vec2 offset = dir * float(i) * texelSize;
-        vec4 sample_color = texture(u_image0, v_texCoord + offset);
-
-        float weight;
-        if (u_int0 == BLUR_GAUSSIAN) {
-            weight = gaussian(float(i), sigma);
-        } else {
-            // BLUR_BOX
-            weight = 1.0;
-        }
-
-        color += sample_color * weight;
-        totalWeight += weight;
-    }
-
-    fragColor0 = color / totalWeight;
-}
@@ -1,19 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-
-in vec2 v_texCoord;
-layout(location = 0) out vec4 fragColor0;
-layout(location = 1) out vec4 fragColor1;
-layout(location = 2) out vec4 fragColor2;
-layout(location = 3) out vec4 fragColor3;
-
-void main() {
-  vec4 color = texture(u_image0, v_texCoord);
-  // Output each channel as grayscale to separate render targets
-  fragColor0 = vec4(vec3(color.r), 1.0);  // Red channel
-  fragColor1 = vec4(vec3(color.g), 1.0);  // Green channel
-  fragColor2 = vec4(vec3(color.b), 1.0);  // Blue channel
-  fragColor3 = vec4(vec3(color.a), 1.0);  // Alpha channel
-}
@@ -1,71 +0,0 @@
-#version 300 es
-precision highp float;
-
-// Levels Adjustment
-// u_int0:   channel      (0=RGB, 1=R, 2=G, 3=B)         default: 0
-// u_float0: input black  (0-255)                        default: 0
-// u_float1: input white  (0-255)                        default: 255
-// u_float2: gamma        (0.01-9.99)                    default: 1.0
-// u_float3: output black (0-255)                        default: 0
-// u_float4: output white (0-255)                        default: 255
-
-uniform sampler2D u_image0;
-uniform int u_int0;
-uniform float u_float0;
-uniform float u_float1;
-uniform float u_float2;
-uniform float u_float3;
-uniform float u_float4;
-
-in vec2 v_texCoord;
-out vec4 fragColor;
-
-vec3 applyLevels(vec3 color, float inBlack, float inWhite, float gamma, float outBlack, float outWhite) {
-    float inRange = max(inWhite - inBlack, 0.0001);
-    vec3 result = clamp((color - inBlack) / inRange, 0.0, 1.0);
-    result = pow(result, vec3(1.0 / gamma));
-    result = mix(vec3(outBlack), vec3(outWhite), result);
-    return result;
-}
-
-float applySingleChannel(float value, float inBlack, float inWhite, float gamma, float outBlack, float outWhite) {
-    float inRange = max(inWhite - inBlack, 0.0001);
-    float result = clamp((value - inBlack) / inRange, 0.0, 1.0);
-    result = pow(result, 1.0 / gamma);
-    result = mix(outBlack, outWhite, result);
-    return result;
-}
-
-void main() {
-    vec4 texColor = texture(u_image0, v_texCoord);
-    vec3 color = texColor.rgb;
-    
-    float inBlack = u_float0 / 255.0;
-    float inWhite = u_float1 / 255.0;
-    float gamma = u_float2;
-    float outBlack = u_float3 / 255.0;
-    float outWhite = u_float4 / 255.0;
-    
-    vec3 result;
-    
-    if (u_int0 == 0) {
-        result = applyLevels(color, inBlack, inWhite, gamma, outBlack, outWhite);
-    }
-    else if (u_int0 == 1) {
-        result = color;
-        result.r = applySingleChannel(color.r, inBlack, inWhite, gamma, outBlack, outWhite);
-    }
-    else if (u_int0 == 2) {
-        result = color;
-        result.g = applySingleChannel(color.g, inBlack, inWhite, gamma, outBlack, outWhite);
-    }
-    else if (u_int0 == 3) {
-        result = color;
-        result.b = applySingleChannel(color.b, inBlack, inWhite, gamma, outBlack, outWhite);
-    }
-    else {
-        result = color;
-    }
-    
-    fragColor = vec4(result, texColor.a);
-}
@@ -1,28 +0,0 @@
-# GLSL Shader Sources
-
-This folder contains the GLSL fragment shaders extracted from blueprint JSON files for easier editing and version control.
-
-## File Naming Convention
-
-`{Blueprint_Name}_{node_id}.frag`
-
- **Blueprint_Name**: The JSON filename with spaces/special chars replaced by underscores
- **node_id**: The GLSLShader node ID within the subgraph
-
-## Usage
-
-```bash
-# Extract shaders from blueprint JSONs to this folder
-python update_blueprints.py extract
-
-# Patch edited shaders back into blueprint JSONs
-python update_blueprints.py patch
-```
-
-## Workflow
-
-1. Run `extract` to pull current shaders from JSONs
-2. Edit `.frag` files
-3. Run `patch` to update the blueprint JSONs
-4. Test
-5. Commit both `.frag` files and updated JSONs
@@ -1,28 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform vec2 u_resolution;
-uniform float u_float0;  // strength [0.0 – 2.0] typical: 0.3–1.0
-
-in vec2 v_texCoord;
-layout(location = 0) out vec4 fragColor0;
-
-void main() {
-    vec2 texel = 1.0 / u_resolution;
-    
-    // Sample center and neighbors
-    vec4 center = texture(u_image0, v_texCoord);
-    vec4 top    = texture(u_image0, v_texCoord + vec2( 0.0, -texel.y));
-    vec4 bottom = texture(u_image0, v_texCoord + vec2( 0.0,  texel.y));
-    vec4 left   = texture(u_image0, v_texCoord + vec2(-texel.x,  0.0));
-    vec4 right  = texture(u_image0, v_texCoord + vec2( texel.x,  0.0));
-    
-    // Edge enhancement (Laplacian)
-    vec4 edges = center * 4.0 - top - bottom - left - right;
-    
-    // Add edges back scaled by strength
-    vec4 sharpened = center + edges * u_float0;
-    
-    fragColor0 = vec4(clamp(sharpened.rgb, 0.0, 1.0), center.a);
-}
@@ -1,61 +0,0 @@
-#version 300 es
-precision highp float;
-
-uniform sampler2D u_image0;
-uniform vec2 u_resolution;
-uniform float u_float0;  // amount    [0.0 - 3.0]  typical: 0.5-1.5
-uniform float u_float1;  // radius    [0.5 - 10.0] blur radius in pixels
-uniform float u_float2;  // threshold [0.0 - 0.1]  min difference to sharpen
-
-in vec2 v_texCoord;
-layout(location = 0) out vec4 fragColor0;
-
-float gaussian(float x, float sigma) {
-    return exp(-(x * x) / (2.0 * sigma * sigma));
-}
-
-float getLuminance(vec3 color) {
-    return dot(color, vec3(0.2126, 0.7152, 0.0722));
-}
-
-void main() {
-    vec2 texel = 1.0 / u_resolution;
-    float radius = max(u_float1, 0.5);
-    float amount = u_float0;
-    float threshold = u_float2;
-
-    vec4 original = texture(u_image0, v_texCoord);
-
-    // Gaussian blur for the "unsharp" mask
-    int samples = int(ceil(radius));
-    float sigma = radius / 2.0;
-
-    vec4 blurred = vec4(0.0);
-    float totalWeight = 0.0;
-
-    for (int x = -samples; x <= samples; x++) {
-        for (int y = -samples; y <= samples; y++) {
-            vec2 offset = vec2(float(x), float(y)) * texel;
-            vec4 sample_color = texture(u_image0, v_texCoord + offset);
-
-            float dist = length(vec2(float(x), float(y)));
-            float weight = gaussian(dist, sigma);
-            blurred += sample_color * weight;
-            totalWeight += weight;
-        }
-    }
-    blurred /= totalWeight;
-
-    // Unsharp mask = original - blurred
-    vec3 mask = original.rgb - blurred.rgb;
-
-    // Luminance-based threshold with smooth falloff
-    float lumaDelta = abs(getLuminance(original.rgb) - getLuminance(blurred.rgb));
-    float thresholdScale = smoothstep(0.0, threshold, lumaDelta);
-    mask *= thresholdScale;
-
-    // Sharpen: original + mask * amount
-    vec3 sharpened = original.rgb + mask * amount;
-
-    fragColor0 = vec4(clamp(sharpened, 0.0, 1.0), original.a);
-}
@@ -1,159 +0,0 @@
-#!/usr/bin/env python3
-"""
-Shader Blueprint Updater
-
-Syncs GLSL shader files between this folder and blueprint JSON files.
-
-File naming convention:
-    {Blueprint Name}_{node_id}.frag
-
-Usage:
-    python update_blueprints.py extract   # Extract shaders from JSONs to here
-    python update_blueprints.py patch     # Patch shaders back into JSONs
-    python update_blueprints.py           # Same as patch (default)
-"""
-
-import json
-import logging
-import sys
-import re
-from pathlib import Path
-
-logging.basicConfig(level=logging.INFO, format='%(message)s')
-logger = logging.getLogger(__name__)
-
-GLSL_DIR = Path(__file__).parent
-BLUEPRINTS_DIR = GLSL_DIR.parent
-
-
-def get_blueprint_files():
-    """Get all blueprint JSON files."""
-    return sorted(BLUEPRINTS_DIR.glob("*.json"))
-
-
-def sanitize_filename(name):
-    """Convert blueprint name to safe filename."""
-    return re.sub(r'[^\w\-]', '_', name)
-
-
-def extract_shaders():
-    """Extract all shaders from blueprint JSONs to this folder."""
-    extracted = 0
-    for json_path in get_blueprint_files():
-        blueprint_name = json_path.stem
-
-        try:
-            with open(json_path, 'r') as f:
-                data = json.load(f)
-        except (json.JSONDecodeError, IOError) as e:
-            logger.warning("Skipping %s: %s", json_path.name, e)
-            continue
-
-        # Find GLSLShader nodes in subgraphs
-        for subgraph in data.get('definitions', {}).get('subgraphs', []):
-            for node in subgraph.get('nodes', []):
-                if node.get('type') == 'GLSLShader':
-                    node_id = node.get('id')
-                    widgets = node.get('widgets_values', [])
-
-                    # Find shader code (first string that looks like GLSL)
-                    for widget in widgets:
-                        if isinstance(widget, str) and widget.startswith('#version'):
-                            safe_name = sanitize_filename(blueprint_name)
-                            frag_name = f"{safe_name}_{node_id}.frag"
-                            frag_path = GLSL_DIR / frag_name
-
-                            with open(frag_path, 'w') as f:
-                                f.write(widget)
-
-                            logger.info("  Extracted: %s", frag_name)
-                            extracted += 1
-                            break
-
-    logger.info("\nExtracted %d shader(s)", extracted)
-
-
-def patch_shaders():
-    """Patch shaders from this folder back into blueprint JSONs."""
-    # Build lookup: blueprint_name -> [(node_id, shader_code), ...]
-    shader_updates = {}
-
-    for frag_path in sorted(GLSL_DIR.glob("*.frag")):
-        # Parse filename: {blueprint_name}_{node_id}.frag
-        parts = frag_path.stem.rsplit('_', 1)
-        if len(parts) != 2:
-            logger.warning("Skipping %s: invalid filename format", frag_path.name)
-            continue
-
-        blueprint_name, node_id_str = parts
-
-        try:
-            node_id = int(node_id_str)
-        except ValueError:
-            logger.warning("Skipping %s: invalid node_id", frag_path.name)
-            continue
-
-        with open(frag_path, 'r') as f:
-            shader_code = f.read()
-
-        if blueprint_name not in shader_updates:
-            shader_updates[blueprint_name] = []
-        shader_updates[blueprint_name].append((node_id, shader_code))
-
-    # Apply updates to JSON files
-    patched = 0
-    for json_path in get_blueprint_files():
-        blueprint_name = sanitize_filename(json_path.stem)
-
-        if blueprint_name not in shader_updates:
-            continue
-
-        try:
-            with open(json_path, 'r') as f:
-                data = json.load(f)
-        except (json.JSONDecodeError, IOError) as e:
-            logger.error("Error reading %s: %s", json_path.name, e)
-            continue
-
-        modified = False
-        for node_id, shader_code in shader_updates[blueprint_name]:
-            # Find the node and update
-            for subgraph in data.get('definitions', {}).get('subgraphs', []):
-                for node in subgraph.get('nodes', []):
-                    if node.get('id') == node_id and node.get('type') == 'GLSLShader':
-                        widgets = node.get('widgets_values', [])
-                        if len(widgets) > 0 and widgets[0] != shader_code:
-                            widgets[0] = shader_code
-                            modified = True
-                            logger.info("  Patched: %s (node %d)", json_path.name, node_id)
-                            patched += 1
-
-        if modified:
-            with open(json_path, 'w') as f:
-                json.dump(data, f)
-
-    if patched == 0:
-        logger.info("No changes to apply.")
-    else:
-        logger.info("\nPatched %d shader(s)", patched)
-
-
-def main():
-    if len(sys.argv) < 2:
-        command = "patch"
-    else:
-        command = sys.argv[1].lower()
-
-    if command == "extract":
-        logger.info("Extracting shaders from blueprints...")
-        extract_shaders()
-    elif command in ("patch", "update", "apply"):
-        logger.info("Patching shaders into blueprints...")
-        patch_shaders()
-    else:
-        logger.info(__doc__)
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()
@@ -1 +0,0 @@
-{"revision": 0, "last_node_id": 29, "last_link_id": 0, "nodes": [{"id": 29, "type": "4c9d6ea4-b912-40e5-8766-6793a9758c53", "pos": [1970, -230], "size": [180, 86], "flags": {}, "order": 5, "mode": 0, "inputs": [{"label": "image", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": null}], "outputs": [{"label": "R", "localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": []}, {"label": "G", "localized_name": "IMAGE1", "name": "IMAGE1", "type": "IMAGE", "links": []}, {"label": "B", "localized_name": "IMAGE2", "name": "IMAGE2", "type": "IMAGE", "links": []}, {"label": "A", "localized_name": "IMAGE3", "name": "IMAGE3", "type": "IMAGE", "links": []}], "title": "Image Channels", "properties": {"proxyWidgets": []}, "widgets_values": []}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "4c9d6ea4-b912-40e5-8766-6793a9758c53", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 28, "lastLinkId": 39, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Image Channels", "inputNode": {"id": -10, "bounding": [1820, -185, 120, 60]}, "outputNode": {"id": -20, "bounding": [2460, -215, 120, 120]}, "inputs": [{"id": "3522932b-2d86-4a1f-a02a-cb29f3a9d7fe", "name": "images.image0", "type": "IMAGE", "linkIds": [39], "localized_name": "images.image0", "label": "image", "pos": [1920, -165]}], "outputs": [{"id": "605cb9c3-b065-4d9b-81d2-3ec331889b2b", "name": "IMAGE0", "type": "IMAGE", "linkIds": [26], "localized_name": "IMAGE0", "label": "R", "pos": [2480, -195]}, {"id": "fb44a77e-0522-43e9-9527-82e7465b3596", "name": "IMAGE1", "type": "IMAGE", "linkIds": [27], "localized_name": "IMAGE1", "label": "G", "pos": [2480, -175]}, {"id": "81460ee6-0131-402a-874f-6bf3001fc4ff", "name": "IMAGE2", "type": "IMAGE", "linkIds": [28], "localized_name": "IMAGE2", "label": "B", "pos": [2480, -155]}, {"id": "ae690246-80d4-4951-b1d9-9306d8a77417", "name": "IMAGE3", "type": "IMAGE", "linkIds": [29], "localized_name": "IMAGE3", "label": "A", "pos": [2480, -135]}], "widgets": [], "nodes": [{"id": 23, "type": "GLSLShader", "pos": [2000, -330], "size": [400, 172], "flags": {}, "order": 0, "mode": 0, "inputs": [{"label": "image", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": 39}, {"localized_name": "fragment_shader", "name": "fragment_shader", "type": "STRING", "widget": {"name": "fragment_shader"}, "link": null}, {"localized_name": "size_mode", "name": "size_mode", "type": "COMFY_DYNAMICCOMBO_V3", "widget": {"name": "size_mode"}, "link": null}, {"label": "image1", "localized_name": "images.image1", "name": "images.image1", "shape": 7, "type": "IMAGE", "link": null}], "outputs": [{"label": "R", "localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": [26]}, {"label": "G", "localized_name": "IMAGE1", "name": "IMAGE1", "type": "IMAGE", "links": [27]}, {"label": "B", "localized_name": "IMAGE2", "name": "IMAGE2", "type": "IMAGE", "links": [28]}, {"label": "A", "localized_name": "IMAGE3", "name": "IMAGE3", "type": "IMAGE", "links": [29]}], "properties": {"Node name for S&R": "GLSLShader"}, "widgets_values": ["#version 300 es\nprecision highp float;\n\nuniform sampler2D u_image0;\n\nin vec2 v_texCoord;\nlayout(location = 0) out vec4 fragColor0;\nlayout(location = 1) out vec4 fragColor1;\nlayout(location = 2) out vec4 fragColor2;\nlayout(location = 3) out vec4 fragColor3;\n\nvoid main() {\n  vec4 color = texture(u_image0, v_texCoord);\n  // Output each channel as grayscale to separate render targets\n  fragColor0 = vec4(vec3(color.r), 1.0);  // Red channel\n  fragColor1 = vec4(vec3(color.g), 1.0);  // Green channel\n  fragColor2 = vec4(vec3(color.b), 1.0);  // Blue channel\n  fragColor3 = vec4(vec3(color.a), 1.0);  // Alpha channel\n}\n", "from_input"]}], "groups": [], "links": [{"id": 39, "origin_id": -10, "origin_slot": 0, "target_id": 23, "target_slot": 0, "type": "IMAGE"}, {"id": 26, "origin_id": 23, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "IMAGE"}, {"id": 27, "origin_id": 23, "origin_slot": 1, "target_id": -20, "target_slot": 1, "type": "IMAGE"}, {"id": 28, "origin_id": 23, "origin_slot": 2, "target_id": -20, "target_slot": 2, "type": "IMAGE"}, {"id": 29, "origin_id": 23, "origin_slot": 3, "target_id": -20, "target_slot": 3, "type": "IMAGE"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Image Tools/Color adjust"}]}}
@@ -1 +0,0 @@
-{"revision": 0, "last_node_id": 15, "last_link_id": 0, "nodes": [{"id": 15, "type": "24d8bbfd-39d4-4774-bff0-3de40cc7a471", "pos": [-1490, 2040], "size": [400, 260], "flags": {}, "order": 0, "mode": 0, "inputs": [{"name": "prompt", "type": "STRING", "widget": {"name": "prompt"}, "link": null}, {"label": "reference images", "name": "images", "type": "IMAGE", "link": null}], "outputs": [{"name": "STRING", "type": "STRING", "links": null}], "title": "Prompt Enhance", "properties": {"proxyWidgets": [["-1", "prompt"]], "cnr_id": "comfy-core", "ver": "0.14.1"}, "widgets_values": [""]}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "24d8bbfd-39d4-4774-bff0-3de40cc7a471", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 15, "lastLinkId": 14, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Prompt Enhance", "inputNode": {"id": -10, "bounding": [-2170, 2110, 138.876953125, 80]}, "outputNode": {"id": -20, "bounding": [-640, 2110, 120, 60]}, "inputs": [{"id": "aeab7216-00e0-4528-a09b-bba50845c5a6", "name": "prompt", "type": "STRING", "linkIds": [11], "pos": [-2051.123046875, 2130]}, {"id": "7b73fd36-aa31-4771-9066-f6c83879994b", "name": "images", "type": "IMAGE", "linkIds": [14], "label": "reference images", "pos": [-2051.123046875, 2150]}], "outputs": [{"id": "c7b0d930-68a1-48d1-b496-0519e5837064", "name": "STRING", "type": "STRING", "linkIds": [13], "pos": [-620, 2130]}], "widgets": [], "nodes": [{"id": 11, "type": "GeminiNode", "pos": [-1560, 1990], "size": [470, 470], "flags": {}, "order": 0, "mode": 0, "inputs": [{"localized_name": "images", "name": "images", "shape": 7, "type": "IMAGE", "link": 14}, {"localized_name": "audio", "name": "audio", "shape": 7, "type": "AUDIO", "link": null}, {"localized_name": "video", "name": "video", "shape": 7, "type": "VIDEO", "link": null}, {"localized_name": "files", "name": "files", "shape": 7, "type": "GEMINI_INPUT_FILES", "link": null}, {"localized_name": "prompt", "name": "prompt", "type": "STRING", "widget": {"name": "prompt"}, "link": 11}, {"localized_name": "model", "name": "model", "type": "COMBO", "widget": {"name": "model"}, "link": null}, {"localized_name": "seed", "name": "seed", "type": "INT", "widget": {"name": "seed"}, "link": null}, {"localized_name": "system_prompt", "name": "system_prompt", "shape": 7, "type": "STRING", "widget": {"name": "system_prompt"}, "link": null}], "outputs": [{"localized_name": "STRING", "name": "STRING", "type": "STRING", "links": [13]}], "properties": {"cnr_id": "comfy-core", "ver": "0.14.1", "Node name for S&R": "GeminiNode"}, "widgets_values": ["", "gemini-3-pro-preview", 42, "randomize", "You are an expert in prompt writing.\nBased on the input, rewrite the user's input into a detailed prompt.\nincluding camera settings, lighting, composition, and style.\nReturn the prompt only"], "color": "#432", "bgcolor": "#653"}], "groups": [], "links": [{"id": 11, "origin_id": -10, "origin_slot": 0, "target_id": 11, "target_slot": 4, "type": "STRING"}, {"id": 13, "origin_id": 11, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "STRING"}, {"id": 14, "origin_id": -10, "origin_slot": 1, "target_id": 11, "target_slot": 0, "type": "IMAGE"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Text generation/Prompt enhance"}]}, "extra": {}}
@@ -1 +0,0 @@
-{"revision": 0, "last_node_id": 25, "last_link_id": 0, "nodes": [{"id": 25, "type": "621ba4e2-22a8-482d-a369-023753198b7b", "pos": [4610, -790], "size": [230, 58], "flags": {}, "order": 4, "mode": 0, "inputs": [{"label": "image", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": null}], "outputs": [{"label": "IMAGE", "localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": []}], "title": "Sharpen", "properties": {"proxyWidgets": [["24", "value"]]}, "widgets_values": []}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "621ba4e2-22a8-482d-a369-023753198b7b", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 24, "lastLinkId": 36, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Sharpen", "inputNode": {"id": -10, "bounding": [4090, -825, 120, 60]}, "outputNode": {"id": -20, "bounding": [5150, -825, 120, 60]}, "inputs": [{"id": "37011fb7-14b7-4e0e-b1a0-6a02e8da1fd7", "name": "images.image0", "type": "IMAGE", "linkIds": [34], "localized_name": "images.image0", "label": "image", "pos": [4190, -805]}], "outputs": [{"id": "e9182b3f-635c-4cd4-a152-4b4be17ae4b9", "name": "IMAGE0", "type": "IMAGE", "linkIds": [35], "localized_name": "IMAGE0", "label": "IMAGE", "pos": [5170, -805]}], "widgets": [], "nodes": [{"id": 24, "type": "PrimitiveFloat", "pos": [4280, -1240], "size": [270, 58], "flags": {}, "order": 0, "mode": 0, "inputs": [{"label": "strength", "localized_name": "value", "name": "value", "type": "FLOAT", "widget": {"name": "value"}, "link": null}], "outputs": [{"localized_name": "FLOAT", "name": "FLOAT", "type": "FLOAT", "links": [36]}], "properties": {"Node name for S&R": "PrimitiveFloat", "min": 0, "max": 3, "precision": 2, "step": 0.05}, "widgets_values": [0.5]}, {"id": 23, "type": "GLSLShader", "pos": [4570, -1240], "size": [370, 192], "flags": {}, "order": 1, "mode": 0, "inputs": [{"label": "image0", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": 34}, {"label": "image1", "localized_name": "images.image1", "name": "images.image1", "shape": 7, "type": "IMAGE", "link": null}, {"label": "u_float0", "localized_name": "floats.u_float0", "name": "floats.u_float0", "shape": 7, "type": "FLOAT", "link": 36}, {"label": "u_float1", "localized_name": "floats.u_float1", "name": "floats.u_float1", "shape": 7, "type": "FLOAT", "link": null}, {"label": "u_int0", "localized_name": "ints.u_int0", "name": "ints.u_int0", "shape": 7, "type": "INT", "link": null}, {"localized_name": "fragment_shader", "name": "fragment_shader", "type": "STRING", "widget": {"name": "fragment_shader"}, "link": null}, {"localized_name": "size_mode", "name": "size_mode", "type": "COMFY_DYNAMICCOMBO_V3", "widget": {"name": "size_mode"}, "link": null}], "outputs": [{"localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": [35]}, {"localized_name": "IMAGE1", "name": "IMAGE1", "type": "IMAGE", "links": null}, {"localized_name": "IMAGE2", "name": "IMAGE2", "type": "IMAGE", "links": null}, {"localized_name": "IMAGE3", "name": "IMAGE3", "type": "IMAGE", "links": null}], "properties": {"Node name for S&R": "GLSLShader"}, "widgets_values": ["#version 300 es\nprecision highp float;\n\nuniform sampler2D u_image0;\nuniform vec2 u_resolution;\nuniform float u_float0;  // strength [0.0 – 2.0] typical: 0.3–1.0\n\nin vec2 v_texCoord;\nlayout(location = 0) out vec4 fragColor0;\n\nvoid main() {\n    vec2 texel = 1.0 / u_resolution;\n    \n    // Sample center and neighbors\n    vec4 center = texture(u_image0, v_texCoord);\n    vec4 top    = texture(u_image0, v_texCoord + vec2( 0.0, -texel.y));\n    vec4 bottom = texture(u_image0, v_texCoord + vec2( 0.0,  texel.y));\n    vec4 left   = texture(u_image0, v_texCoord + vec2(-texel.x,  0.0));\n    vec4 right  = texture(u_image0, v_texCoord + vec2( texel.x,  0.0));\n    \n    // Edge enhancement (Laplacian)\n    vec4 edges = center * 4.0 - top - bottom - left - right;\n    \n    // Add edges back scaled by strength\n    vec4 sharpened = center + edges * u_float0;\n    \n    fragColor0 = vec4(clamp(sharpened.rgb, 0.0, 1.0), center.a);\n}", "from_input"]}], "groups": [], "links": [{"id": 36, "origin_id": 24, "origin_slot": 0, "target_id": 23, "target_slot": 2, "type": "FLOAT"}, {"id": 34, "origin_id": -10, "origin_slot": 0, "target_id": 23, "target_slot": 0, "type": "IMAGE"}, {"id": 35, "origin_id": 23, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "IMAGE"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Image Tools/Sharpen"}]}}
@@ -1 +0,0 @@
-{"revision": 0, "last_node_id": 13, "last_link_id": 0, "nodes": [{"id": 13, "type": "cf95b747-3e17-46cb-8097-cac60ff9b2e1", "pos": [1120, 330], "size": [240, 58], "flags": {}, "order": 3, "mode": 0, "inputs": [{"localized_name": "video", "name": "video", "type": "VIDEO", "link": null}, {"name": "model_name", "type": "COMBO", "widget": {"name": "model_name"}, "link": null}], "outputs": [{"localized_name": "VIDEO", "name": "VIDEO", "type": "VIDEO", "links": []}], "title": "Video Upscale(GAN x4)", "properties": {"proxyWidgets": [["-1", "model_name"]], "cnr_id": "comfy-core", "ver": "0.14.1"}, "widgets_values": ["RealESRGAN_x4plus.safetensors"]}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "cf95b747-3e17-46cb-8097-cac60ff9b2e1", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 13, "lastLinkId": 19, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Video Upscale(GAN x4)", "inputNode": {"id": -10, "bounding": [550, 460, 120, 80]}, "outputNode": {"id": -20, "bounding": [1490, 460, 120, 60]}, "inputs": [{"id": "666d633e-93e7-42dc-8d11-2b7b99b0f2a6", "name": "video", "type": "VIDEO", "linkIds": [10], "localized_name": "video", "pos": [650, 480]}, {"id": "2e23a087-caa8-4d65-99e6-662761aa905a", "name": "model_name", "type": "COMBO", "linkIds": [19], "pos": [650, 500]}], "outputs": [{"id": "0c1768ea-3ec2-412f-9af6-8e0fa36dae70", "name": "VIDEO", "type": "VIDEO", "linkIds": [15], "localized_name": "VIDEO", "pos": [1510, 480]}], "widgets": [], "nodes": [{"id": 2, "type": "ImageUpscaleWithModel", "pos": [1110, 450], "size": [320, 46], "flags": {}, "order": 1, "mode": 0, "inputs": [{"localized_name": "upscale_model", "name": "upscale_model", "type": "UPSCALE_MODEL", "link": 1}, {"localized_name": "image", "name": "image", "type": "IMAGE", "link": 14}], "outputs": [{"localized_name": "IMAGE", "name": "IMAGE", "type": "IMAGE", "links": [13]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "ImageUpscaleWithModel"}}, {"id": 11, "type": "CreateVideo", "pos": [1110, 550], "size": [320, 78], "flags": {}, "order": 3, "mode": 0, "inputs": [{"localized_name": "images", "name": "images", "type": "IMAGE", "link": 13}, {"localized_name": "audio", "name": "audio", "shape": 7, "type": "AUDIO", "link": 16}, {"localized_name": "fps", "name": "fps", "type": "FLOAT", "widget": {"name": "fps"}, "link": 12}], "outputs": [{"localized_name": "VIDEO", "name": "VIDEO", "type": "VIDEO", "links": [15]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "CreateVideo"}, "widgets_values": [30]}, {"id": 10, "type": "GetVideoComponents", "pos": [1110, 330], "size": [320, 70], "flags": {}, "order": 2, "mode": 0, "inputs": [{"localized_name": "video", "name": "video", "type": "VIDEO", "link": 10}], "outputs": [{"localized_name": "images", "name": "images", "type": "IMAGE", "links": [14]}, {"localized_name": "audio", "name": "audio", "type": "AUDIO", "links": [16]}, {"localized_name": "fps", "name": "fps", "type": "FLOAT", "links": [12]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "GetVideoComponents"}}, {"id": 1, "type": "UpscaleModelLoader", "pos": [750, 450], "size": [280, 60], "flags": {}, "order": 0, "mode": 0, "inputs": [{"localized_name": "model_name", "name": "model_name", "type": "COMBO", "widget": {"name": "model_name"}, "link": 19}], "outputs": [{"localized_name": "UPSCALE_MODEL", "name": "UPSCALE_MODEL", "type": "UPSCALE_MODEL", "links": [1]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "UpscaleModelLoader", "models": [{"name": "RealESRGAN_x4plus.safetensors", "url": "https://huggingface.co/Comfy-Org/Real-ESRGAN_repackaged/resolve/main/RealESRGAN_x4plus.safetensors", "directory": "upscale_models"}]}, "widgets_values": ["RealESRGAN_x4plus.safetensors"]}], "groups": [], "links": [{"id": 1, "origin_id": 1, "origin_slot": 0, "target_id": 2, "target_slot": 0, "type": "UPSCALE_MODEL"}, {"id": 14, "origin_id": 10, "origin_slot": 0, "target_id": 2, "target_slot": 1, "type": "IMAGE"}, {"id": 13, "origin_id": 2, "origin_slot": 0, "target_id": 11, "target_slot": 0, "type": "IMAGE"}, {"id": 16, "origin_id": 10, "origin_slot": 1, "target_id": 11, "target_slot": 1, "type": "AUDIO"}, {"id": 12, "origin_id": 10, "origin_slot": 2, "target_id": 11, "target_slot": 2, "type": "FLOAT"}, {"id": 10, "origin_id": -10, "origin_slot": 0, "target_id": 10, "target_slot": 0, "type": "VIDEO"}, {"id": 15, "origin_id": 11, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "VIDEO"}, {"id": 19, "origin_id": -10, "origin_slot": 1, "target_id": 1, "target_slot": 0, "type": "COMBO"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Video generation and editing/Enhance video"}]}, "extra": {}}
@@ -25,11 +25,11 @@ class AudioEncoderModel():
        elif model_type == "whisper3":
            self.model = WhisperLargeV3(**model_config)
        self.model.eval()
-        self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
+        self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
        self.model_sample_rate = 16000

    def load_sd(self, sd):
-        return self.model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
+        return self.model.load_state_dict(sd, strict=False)

    def get_sd(self):
        return self.model.state_dict()
@@ -0,0 +1,13 @@
+import pickle
+
+load = pickle.load
+
+class Empty:
+    pass
+
+class Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        #TODO: safe unpickle
+        if module.startswith("pytorch_lightning"):
+            return Empty
+        return super().find_class(module, name)
@@ -159,7 +159,6 @@ class PerformanceFeature(enum.Enum):
    Fp8MatrixMultiplication = "fp8_matrix_mult"
    CublasOps = "cublas_ops"
    AutoTune = "autotune"
-    DynamicVRAM = "dynamic_vram"

 parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. This is used to test new features so using it might crash your comfyui. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: {}".format(" ".join(map(lambda c: c.value, PerformanceFeature))))

@@ -258,6 +257,3 @@ elif args.fast == []:
 # '--fast' is provided with a list of performance features, use that list
 else:
    args.fast = set(args.fast)
-
-def enables_dynamic_vram():
-    return PerformanceFeature.DynamicVRAM in args.fast and not args.highvram and not args.gpu_only
@@ -47,10 +47,10 @@ class ClipVisionModel():
        self.model = model_class(config, self.dtype, offload_device, comfy.ops.manual_cast)
        self.model.eval()

-        self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
+        self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)

    def load_sd(self, sd):
-        return self.model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
+        return self.model.load_state_dict(sd, strict=False)

    def get_sd(self):
        return self.model.state_dict()
@@ -176,8 +176,6 @@ class InputTypeOptions(TypedDict):
    """COMBO type only. Specifies the configuration for a multi-select widget.
    Available after ComfyUI frontend v1.13.4
    https://github.com/Comfy-Org/ComfyUI_frontend/pull/2987"""
-    gradient_stops: NotRequired[list[list[float]]]
-    """Gradient color stops for gradientslider display mode. Each stop is [offset, r, g, b] (``FLOAT``)."""


 class HiddenInputTypeDict(TypedDict):
@@ -203,7 +203,7 @@ class ControlNet(ControlBase):
        self.control_model = control_model
        self.load_device = load_device
        if control_model is not None:
-            self.control_model_wrapped = comfy.model_patcher.CoreModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())
+            self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())

        self.compression_ratio = compression_ratio
        self.global_average_pooling = global_average_pooling
@@ -297,30 +297,6 @@ class ControlNet(ControlBase):
        self.model_sampling_current = None
        super().cleanup()

-
-class QwenFunControlNet(ControlNet):
-    def get_control(self, x_noisy, t, cond, batched_number, transformer_options):
-        # Fun checkpoints are more sensitive to high strengths in the generic
-        # ControlNet merge path. Use a soft response curve so strength=1.0 stays
-        # unchanged while >1 grows more gently.
-        original_strength = self.strength
-        self.strength = math.sqrt(max(self.strength, 0.0))
-        try:
-            return super().get_control(x_noisy, t, cond, batched_number, transformer_options)
-        finally:
-            self.strength = original_strength
-
-    def pre_run(self, model, percent_to_timestep_function):
-        super().pre_run(model, percent_to_timestep_function)
-        self.set_extra_arg("base_model", model.diffusion_model)
-
-    def copy(self):
-        c = QwenFunControlNet(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
-        c.control_model = self.control_model
-        c.control_model_wrapped = self.control_model_wrapped
-        self.copy_to(c)
-        return c
-
 class ControlLoraOps:
    class Linear(torch.nn.Module, comfy.ops.CastWeightBiasOp):
        def __init__(self, in_features: int, out_features: int, bias: bool = True,
@@ -584,7 +560,6 @@ def load_controlnet_hunyuandit(controlnet_data, model_options={}):
 def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False, model_options={}):
    model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd, model_options=model_options)
    control_model = comfy.ldm.flux.controlnet.ControlNetFlux(mistoline=mistoline, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
-    sd = model_config.process_unet_state_dict(sd)
    control_model = controlnet_load_state_dict(control_model, sd)
    extra_conds = ['y', 'guidance']
    control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
@@ -630,53 +605,6 @@ def load_controlnet_qwen_instantx(sd, model_options={}):
    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
    return control

-
-def load_controlnet_qwen_fun(sd, model_options={}):
-    load_device = comfy.model_management.get_torch_device()
-    weight_dtype = comfy.utils.weight_dtype(sd)
-    unet_dtype = model_options.get("dtype", weight_dtype)
-    manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
-
-    operations = model_options.get("custom_operations", None)
-    if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)
-
-    in_features = sd["control_img_in.weight"].shape[1]
-    inner_dim = sd["control_img_in.weight"].shape[0]
-
-    block_weight = sd["control_blocks.0.attn.to_q.weight"]
-    attention_head_dim = sd["control_blocks.0.attn.norm_q.weight"].shape[0]
-    num_attention_heads = max(1, block_weight.shape[0] // max(1, attention_head_dim))
-
-    model = comfy.ldm.qwen_image.controlnet.QwenImageFunControlNetModel(
-        control_in_features=in_features,
-        inner_dim=inner_dim,
-        num_attention_heads=num_attention_heads,
-        attention_head_dim=attention_head_dim,
-        num_control_blocks=5,
-        main_model_double=60,
-        injection_layers=(0, 12, 24, 36, 48),
-        operations=operations,
-        device=comfy.model_management.unet_offload_device(),
-        dtype=unet_dtype,
-    )
-    model = controlnet_load_state_dict(model, sd)
-
-    latent_format = comfy.latent_formats.Wan21()
-    control = QwenFunControlNet(
-        model,
-        compression_ratio=1,
-        latent_format=latent_format,
-        # Fun checkpoints already expect their own 33-channel context handling.
-        # Enabling generic concat_mask injects an extra mask channel at apply-time
-        # and breaks the intended fallback packing path.
-        concat_mask=False,
-        load_device=load_device,
-        manual_cast_dtype=manual_cast_dtype,
-        extra_conds=[],
-    )
-    return control
-
 def convert_mistoline(sd):
    return comfy.utils.state_dict_prefix_replace(sd, {"single_controlnet_blocks.": "controlnet_single_blocks."})

@@ -754,8 +682,6 @@ def load_controlnet_state_dict(state_dict, model=None, model_options={}):
            return load_controlnet_qwen_instantx(controlnet_data, model_options=model_options)
        elif "controlnet_x_embedder.weight" in controlnet_data:
            return load_controlnet_flux_instantx(controlnet_data, model_options=model_options)
-    elif "control_blocks.0.after_proj.weight" in controlnet_data and "control_img_in.weight" in controlnet_data:
-        return load_controlnet_qwen_fun(controlnet_data, model_options=model_options)

    elif "controlnet_blocks.0.linear.weight" in controlnet_data: #mistoline flux
        return load_controlnet_flux_xlabs_mistoline(convert_mistoline(controlnet_data), mistoline=True, model_options=model_options)
@@ -5,7 +5,7 @@ from scipy import integrate
 import torch
 from torch import nn
 import torchsde
-from tqdm.auto import tqdm
+from tqdm.auto import trange, tqdm

 from . import utils
 from . import deis
@@ -13,9 +13,6 @@ from . import sa_solver
 import comfy.model_patcher
 import comfy.model_sampling

-import comfy.memory_management
-from comfy.utils import model_trange as trange
-
 def append_zero(x):
    return torch.cat([x, x.new_zeros([1])])

@@ -755,10 +755,6 @@ class ACEAudio(LatentFormat):
    latent_channels = 8
    latent_dimensions = 2

-class ACEAudio15(LatentFormat):
-    latent_channels = 64
-    latent_dimensions = 1
-
 class ChromaRadiance(LatentFormat):
    latent_channels = 3
    spacial_downscale_ratio = 1
@@ -179,8 +179,8 @@ class LLMAdapter(nn.Module):
            if source_attention_mask.ndim == 2:
                source_attention_mask = source_attention_mask.unsqueeze(1).unsqueeze(1)

+        x = self.in_proj(self.embed(target_input_ids))
        context = source_hidden_states
-        x = self.in_proj(self.embed(target_input_ids, out_dtype=context.dtype))
        position_ids = torch.arange(x.shape[1], device=x.device).unsqueeze(0)
        position_ids_context = torch.arange(context.shape[1], device=x.device).unsqueeze(0)
        position_embeddings = self.rotary_emb(x, position_ids)
@@ -195,20 +195,8 @@ class Anima(MiniTrainDIT):
        super().__init__(*args, **kwargs)
        self.llm_adapter = LLMAdapter(device=kwargs.get("device"), dtype=kwargs.get("dtype"), operations=kwargs.get("operations"))

-    def preprocess_text_embeds(self, text_embeds, text_ids, t5xxl_weights=None):
+    def preprocess_text_embeds(self, text_embeds, text_ids):
        if text_ids is not None:
-            out = self.llm_adapter(text_embeds, text_ids)
-            if t5xxl_weights is not None:
-                out = out * t5xxl_weights
-
-            if out.shape[1] < 512:
-                out = torch.nn.functional.pad(out, (0, 0, 0, 512 - out.shape[1]))
-            return out
+            return self.llm_adapter(text_embeds, text_ids)
        else:
            return text_embeds
-
-    def forward(self, x, timesteps, context, **kwargs):
-        t5xxl_ids = kwargs.pop("t5xxl_ids", None)
-        if t5xxl_ids is not None:
-            context = self.preprocess_text_embeds(context, t5xxl_ids, t5xxl_weights=kwargs.pop("t5xxl_weights", None))
-        return super().forward(x, timesteps, context, **kwargs)
@@ -3,6 +3,7 @@ from torch import Tensor, nn

 from comfy.ldm.flux.layers import (
    MLPEmbedder,
+    RMSNorm,
    ModulationOut,
 )

@@ -28,7 +29,7 @@ class Approximator(nn.Module):
        super().__init__()
        self.in_proj = operations.Linear(in_dim, hidden_dim, bias=True, dtype=dtype, device=device)
        self.layers = nn.ModuleList([MLPEmbedder(hidden_dim, hidden_dim, dtype=dtype, device=device, operations=operations) for x in range( n_layers)])
-        self.norms = nn.ModuleList([operations.RMSNorm(hidden_dim, dtype=dtype, device=device) for x in range( n_layers)])
+        self.norms = nn.ModuleList([RMSNorm(hidden_dim, dtype=dtype, device=device, operations=operations) for x in range( n_layers)])
        self.out_proj = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device)

    @property
@@ -152,7 +152,6 @@ class Chroma(nn.Module):
        transformer_options={},
        attn_mask: Tensor = None,
    ) -> Tensor:
-        transformer_options = transformer_options.copy()
        patches_replace = transformer_options.get("patches_replace", {})

        # running on sequences img
@@ -229,7 +228,6 @@ class Chroma(nn.Module):

        transformer_options["total_blocks"] = len(self.single_blocks)
        transformer_options["block_type"] = "single"
-        transformer_options["img_slice"] = [txt.shape[1], img.shape[1]]
        for i, block in enumerate(self.single_blocks):
            transformer_options["block_index"] = i
            if i not in self.skip_dit:
@@ -4,6 +4,8 @@ from functools import lru_cache
 import torch
 from torch import nn

+from comfy.ldm.flux.layers import RMSNorm
+

 class NerfEmbedder(nn.Module):
    """
@@ -143,7 +145,7 @@ class NerfGLUBlock(nn.Module):
        # We now need to generate parameters for 3 matrices.
        total_params = 3 * hidden_size_x**2 * mlp_ratio
        self.param_generator = operations.Linear(hidden_size_s, total_params, dtype=dtype, device=device)
-        self.norm = operations.RMSNorm(hidden_size_x, dtype=dtype, device=device)
+        self.norm = RMSNorm(hidden_size_x, dtype=dtype, device=device, operations=operations)
        self.mlp_ratio = mlp_ratio


@@ -176,7 +178,7 @@ class NerfGLUBlock(nn.Module):
 class NerfFinalLayer(nn.Module):
    def __init__(self, hidden_size, out_channels, dtype=None, device=None, operations=None):
        super().__init__()
-        self.norm = operations.RMSNorm(hidden_size, dtype=dtype, device=device)
+        self.norm = RMSNorm(hidden_size, dtype=dtype, device=device, operations=operations)
        self.linear = operations.Linear(hidden_size, out_channels, dtype=dtype, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -188,7 +190,7 @@ class NerfFinalLayer(nn.Module):
 class NerfFinalLayerConv(nn.Module):
    def __init__(self, hidden_size: int, out_channels: int, dtype=None, device=None, operations=None):
        super().__init__()
-        self.norm = operations.RMSNorm(hidden_size, dtype=dtype, device=device)
+        self.norm = RMSNorm(hidden_size, dtype=dtype, device=device, operations=operations)
        self.conv = operations.Conv2d(
            in_channels=hidden_size,
            out_channels=out_channels,
@@ -335,7 +335,7 @@ class FinalLayer(nn.Module):
        device=None, dtype=None, operations=None
    ):
        super().__init__()
-        self.layer_norm = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.layer_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = operations.Linear(
            hidden_size, spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels, bias=False, device=device, dtype=dtype
        )
@@ -463,8 +463,6 @@ class Block(nn.Module):
        extra_per_block_pos_emb: Optional[torch.Tensor] = None,
        transformer_options: Optional[dict] = {},
    ) -> torch.Tensor:
-        residual_dtype = x_B_T_H_W_D.dtype
-        compute_dtype = emb_B_T_D.dtype
        if extra_per_block_pos_emb is not None:
            x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb

@@ -514,7 +512,7 @@ class Block(nn.Module):
        result_B_T_H_W_D = rearrange(
            self.self_attn(
                # normalized_x_B_T_HW_D,
-                rearrange(normalized_x_B_T_H_W_D.to(compute_dtype), "b t h w d -> b (t h w) d"),
+                rearrange(normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
                None,
                rope_emb=rope_emb_L_1_1_D,
                transformer_options=transformer_options,
@@ -524,7 +522,7 @@ class Block(nn.Module):
            h=H,
            w=W,
        )
-        x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D.to(residual_dtype) * result_B_T_H_W_D.to(residual_dtype)
+        x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D * result_B_T_H_W_D

        def _x_fn(
            _x_B_T_H_W_D: torch.Tensor,
@@ -538,7 +536,7 @@ class Block(nn.Module):
            )
            _result_B_T_H_W_D = rearrange(
                self.cross_attn(
-                    rearrange(_normalized_x_B_T_H_W_D.to(compute_dtype), "b t h w d -> b (t h w) d"),
+                    rearrange(_normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
                    crossattn_emb,
                    rope_emb=rope_emb_L_1_1_D,
                    transformer_options=transformer_options,
@@ -557,7 +555,7 @@ class Block(nn.Module):
            shift_cross_attn_B_T_1_1_D,
            transformer_options=transformer_options,
        )
-        x_B_T_H_W_D = result_B_T_H_W_D.to(residual_dtype) * gate_cross_attn_B_T_1_1_D.to(residual_dtype) + x_B_T_H_W_D
+        x_B_T_H_W_D = result_B_T_H_W_D * gate_cross_attn_B_T_1_1_D + x_B_T_H_W_D

        normalized_x_B_T_H_W_D = _fn(
            x_B_T_H_W_D,
@@ -565,8 +563,8 @@ class Block(nn.Module):
            scale_mlp_B_T_1_1_D,
            shift_mlp_B_T_1_1_D,
        )
-        result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.to(compute_dtype))
-        x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp_B_T_1_1_D.to(residual_dtype) * result_B_T_H_W_D.to(residual_dtype)
+        result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
+        x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp_B_T_1_1_D * result_B_T_H_W_D
        return x_B_T_H_W_D


@@ -878,14 +876,6 @@ class MiniTrainDIT(nn.Module):
            "extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
            "transformer_options": kwargs.get("transformer_options", {}),
        }
-
-        # The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream
-        # in fp32, but run attention and MLP modules in fp16.
-        # An alternate method that clamps fp16 values "works" in the sense that it makes coherent images, but there is noticeable
-        # quality degradation and visual artifacts.
-        if x_B_T_H_W_D.dtype == torch.float16:
-            x_B_T_H_W_D = x_B_T_H_W_D.float()
-
        for block in self.blocks:
            x_B_T_H_W_D = block(
                x_B_T_H_W_D,
@@ -894,6 +884,6 @@ class MiniTrainDIT(nn.Module):
                **block_kwargs,
            )

-        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D.to(crossattn_emb.dtype), t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
        x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)[:, :, :orig_shape[-3], :orig_shape[-2], :orig_shape[-1]]
        return x_B_C_Tt_Hp_Wp
@@ -5,9 +5,9 @@ import torch
 from torch import Tensor, nn

 from .math import attention, rope
+import comfy.ops
+import comfy.ldm.common_dit

-# Fix import for some custom nodes, TODO: delete eventually.
-RMSNorm = None

 class EmbedND(nn.Module):
    def __init__(self, dim: int, theta: int, axes_dim: list):
@@ -87,12 +87,20 @@ def build_mlp(hidden_size, mlp_hidden_dim, mlp_silu_act=False, yak_mlp=False, dt
            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
        )

+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.scale = nn.Parameter(torch.empty((dim), dtype=dtype, device=device))
+
+    def forward(self, x: Tensor):
+        return comfy.ldm.common_dit.rms_norm(x, self.scale, 1e-6)
+

 class QKNorm(torch.nn.Module):
    def __init__(self, dim: int, dtype=None, device=None, operations=None):
        super().__init__()
-        self.query_norm = operations.RMSNorm(dim, dtype=dtype, device=device)
-        self.key_norm = operations.RMSNorm(dim, dtype=dtype, device=device)
+        self.query_norm = RMSNorm(dim, dtype=dtype, device=device, operations=operations)
+        self.key_norm = RMSNorm(dim, dtype=dtype, device=device, operations=operations)

    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple:
        q = self.query_norm(q)
@@ -161,7 +169,7 @@ class SiLUActivation(nn.Module):


 class DoubleStreamBlock(nn.Module):
-    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, modulation=True, mlp_silu_act=False, proj_bias=True, yak_mlp=False, dtype=None, device=None, operations=None):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, flipped_img_txt=False, modulation=True, mlp_silu_act=False, proj_bias=True, yak_mlp=False, dtype=None, device=None, operations=None):
        super().__init__()

        mlp_hidden_dim = int(hidden_size * mlp_ratio)
@@ -189,6 +197,8 @@ class DoubleStreamBlock(nn.Module):

        self.txt_mlp = build_mlp(hidden_size, mlp_hidden_dim, mlp_silu_act=mlp_silu_act, yak_mlp=yak_mlp, dtype=dtype, device=device, operations=operations)

+        self.flipped_img_txt = flipped_img_txt
+
    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims_img=None, modulation_dims_txt=None, transformer_options={}):
        if self.modulation:
            img_mod1, img_mod2 = self.img_mod(vec)
@@ -196,9 +206,6 @@ class DoubleStreamBlock(nn.Module):
        else:
            (img_mod1, img_mod2), (txt_mod1, txt_mod2) = vec

-        transformer_patches = transformer_options.get("patches", {})
-        extra_options = transformer_options.copy()
-
        # prepare image for attention
        img_modulated = self.img_norm1(img)
        img_modulated = apply_mod(img_modulated, (1 + img_mod1.scale), img_mod1.shift, modulation_dims_img)
@@ -217,23 +224,32 @@ class DoubleStreamBlock(nn.Module):
        del txt_qkv
        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)

-        q = torch.cat((txt_q, img_q), dim=2)
-        del txt_q, img_q
-        k = torch.cat((txt_k, img_k), dim=2)
-        del txt_k, img_k
-        v = torch.cat((txt_v, img_v), dim=2)
-        del txt_v, img_v
-        # run actual attention
-        attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
-        del q, k, v
+        if self.flipped_img_txt:
+            q = torch.cat((img_q, txt_q), dim=2)
+            del img_q, txt_q
+            k = torch.cat((img_k, txt_k), dim=2)
+            del img_k, txt_k
+            v = torch.cat((img_v, txt_v), dim=2)
+            del img_v, txt_v
+            # run actual attention
+            attn = attention(q, k, v,
+                             pe=pe, mask=attn_mask, transformer_options=transformer_options)
+            del q, k, v

-        if "attn1_output_patch" in transformer_patches:
-            extra_options["img_slice"] = [txt.shape[1], attn.shape[1]]
-            patch = transformer_patches["attn1_output_patch"]
-            for p in patch:
-                attn = p(attn, extra_options)
+            img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1]:]
+        else:
+            q = torch.cat((txt_q, img_q), dim=2)
+            del txt_q, img_q
+            k = torch.cat((txt_k, img_k), dim=2)
+            del txt_k, img_k
+            v = torch.cat((txt_v, img_v), dim=2)
+            del txt_v, img_v
+            # run actual attention
+            attn = attention(q, k, v,
+                             pe=pe, mask=attn_mask, transformer_options=transformer_options)
+            del q, k, v

-        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1]:]
+            txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1]:]

        # calculate the img bloks
        img += apply_mod(self.img_attn.proj(img_attn), img_mod1.gate, None, modulation_dims_img)
@@ -312,9 +328,6 @@ class SingleStreamBlock(nn.Module):
        else:
            mod = vec

-        transformer_patches = transformer_options.get("patches", {})
-        extra_options = transformer_options.copy()
-
        qkv, mlp = torch.split(self.linear1(apply_mod(self.pre_norm(x), (1 + mod.scale), mod.shift, modulation_dims)), [3 * self.hidden_size, self.mlp_hidden_dim_first], dim=-1)

        q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
@@ -324,12 +337,6 @@ class SingleStreamBlock(nn.Module):
        # compute attention
        attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
        del q, k, v
-
-        if "attn1_output_patch" in transformer_patches:
-            patch = transformer_patches["attn1_output_patch"]
-            for p in patch:
-                attn = p(attn, extra_options)
-
        # compute activation in mlp stream, cat again and run second linear layer
        if self.yak_mlp:
            mlp = self.mlp_act(mlp[..., self.mlp_hidden_dim_first // 2:]) * mlp[..., :self.mlp_hidden_dim_first // 2]
@@ -29,34 +29,19 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
    return out.to(dtype=torch.float32, device=pos.device)


-def _apply_rope1(x: Tensor, freqs_cis: Tensor):
-    x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
-
-    x_out = freqs_cis[..., 0] * x_[..., 0]
-    x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])
-
-    return x_out.reshape(*x.shape).type_as(x)
-
-
-def _apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
-    return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis)
-
-
 try:
    import comfy.quant_ops
-    q_apply_rope = comfy.quant_ops.ck.apply_rope
-    q_apply_rope1 = comfy.quant_ops.ck.apply_rope1
-    def apply_rope(xq, xk, freqs_cis):
-        if comfy.model_management.in_training:
-            return _apply_rope(xq, xk, freqs_cis)
-        else:
-            return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis)
-    def apply_rope1(x, freqs_cis):
-        if comfy.model_management.in_training:
-            return _apply_rope1(x, freqs_cis)
-        else:
-            return q_apply_rope1(x, freqs_cis)
+    apply_rope = comfy.quant_ops.ck.apply_rope
+    apply_rope1 = comfy.quant_ops.ck.apply_rope1
 except:
    logging.warning("No comfy kitchen, using old apply_rope functions.")
-    apply_rope = _apply_rope
-    apply_rope1 = _apply_rope1
+    def apply_rope1(x: Tensor, freqs_cis: Tensor):
+        x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
+
+        x_out = freqs_cis[..., 0] * x_[..., 0]
+        x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])
+
+        return x_out.reshape(*x.shape).type_as(x)
+
+    def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
+        return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis)
@@ -16,6 +16,7 @@ from .layers import (
    SingleStreamBlock,
    timestep_embedding,
    Modulation,
+    RMSNorm
 )

@dataclass
@@ -80,7 +81,7 @@ class Flux(nn.Module):
        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device)

        if params.txt_norm:
-            self.txt_norm = operations.RMSNorm(params.context_in_dim, dtype=dtype, device=device)
+            self.txt_norm = RMSNorm(params.context_in_dim, dtype=dtype, device=device, operations=operations)
        else:
            self.txt_norm = None

@@ -142,7 +143,6 @@ class Flux(nn.Module):
        attn_mask: Tensor = None,
    ) -> Tensor:

-        transformer_options = transformer_options.copy()
        patches = transformer_options.get("patches", {})
        patches_replace = transformer_options.get("patches_replace", {})
        if img.ndim != 3 or txt.ndim != 3:
@@ -232,7 +232,6 @@ class Flux(nn.Module):

        transformer_options["total_blocks"] = len(self.single_blocks)
        transformer_options["block_type"] = "single"
-        transformer_options["img_slice"] = [txt.shape[1], img.shape[1]]
        for i, block in enumerate(self.single_blocks):
            transformer_options["block_index"] = i
            if ("single_block", i) in blocks_replace:
@@ -241,6 +241,7 @@ class HunyuanVideo(nn.Module):
                    self.num_heads,
                    mlp_ratio=params.mlp_ratio,
                    qkv_bias=params.qkv_bias,
+                    flipped_img_txt=True,
                    dtype=dtype, device=device, operations=operations
                )
                for _ in range(params.depth)
@@ -304,7 +305,6 @@ class HunyuanVideo(nn.Module):
        control=None,
        transformer_options={},
    ) -> Tensor:
-        transformer_options = transformer_options.copy()
        patches_replace = transformer_options.get("patches_replace", {})

        initial_shape = list(img.shape)
@@ -378,14 +378,14 @@ class HunyuanVideo(nn.Module):
            extra_txt_ids = torch.zeros((txt_ids.shape[0], txt_vision_states.shape[1], txt_ids.shape[-1]), device=txt_ids.device, dtype=txt_ids.dtype)
            txt_ids = torch.cat((txt_ids, extra_txt_ids), dim=1)

-        ids = torch.cat((txt_ids, img_ids), dim=1)
+        ids = torch.cat((img_ids, txt_ids), dim=1)
        pe = self.pe_embedder(ids)

        img_len = img.shape[1]
        if txt_mask is not None:
            attn_mask_len = img_len + txt.shape[1]
            attn_mask = torch.zeros((1, 1, attn_mask_len), dtype=img.dtype, device=img.device)
-            attn_mask[:, 0, :txt.shape[1]] = txt_mask
+            attn_mask[:, 0, img_len:] = txt_mask
        else:
            attn_mask = None

@@ -413,11 +413,10 @@ class HunyuanVideo(nn.Module):
                    if add is not None:
                        img += add

-        img = torch.cat((txt, img), 1)
+        img = torch.cat((img, txt), 1)

        transformer_options["total_blocks"] = len(self.single_blocks)
        transformer_options["block_type"] = "single"
-        transformer_options["img_slice"] = [txt.shape[1], img.shape[1]]
        for i, block in enumerate(self.single_blocks):
            transformer_options["block_index"] = i
            if ("single_block", i) in blocks_replace:
@@ -436,9 +435,9 @@ class HunyuanVideo(nn.Module):
                if i < len(control_o):
                    add = control_o[i]
                    if add is not None:
-                        img[:, txt.shape[1]: img_len + txt.shape[1]] += add
+                        img[:, : img_len] += add

-        img = img[:, txt.shape[1]: img_len + txt.shape[1]]
+        img = img[:, : img_len]
        if ref_latent is not None:
            img = img[:, ref_latent.shape[1]:]

@@ -109,10 +109,10 @@ class HunyuanVideo15SRModel():
        self.model_class = UPSAMPLERS.get(model_type)
        self.model = self.model_class(**config).eval()

-        self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
+        self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)

    def load_sd(self, sd):
-        return self.model.load_state_dict(sd, strict=True, assign=self.patcher.is_dynamic())
+        return self.model.load_state_dict(sd, strict=True)

    def get_sd(self):
        return self.model.state_dict()
@@ -524,9 +524,6 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha

@wrap_attn
 def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
-    if kwargs.get("low_precision_attention", True) is False:
-        return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
-
    exception_fallback = False
    if skip_reshape:
        b, _, _, dim_head = q.shape
@@ -102,7 +102,19 @@ class VideoConv3d(nn.Module):
        return self.conv(x)

 def interpolate_up(x, scale_factor):
-    return torch.nn.functional.interpolate(x, scale_factor=scale_factor, mode="nearest")
+    try:
+        return torch.nn.functional.interpolate(x, scale_factor=scale_factor, mode="nearest")
+    except: #operation not implemented for bf16
+        orig_shape = list(x.shape)
+        out_shape = orig_shape[:2]
+        for i in range(len(orig_shape) - 2):
+            out_shape.append(round(orig_shape[i + 2] * scale_factor[i]))
+        out = torch.empty(out_shape, dtype=x.dtype, layout=x.layout, device=x.device)
+        split = 8
+        l = out.shape[1] // split
+        for i in range(0, out.shape[1], l):
+            out[:,i:i+l] = torch.nn.functional.interpolate(x[:,i:i+l].to(torch.float32), scale_factor=scale_factor, mode="nearest").to(x.dtype)
+        return out

 class Upsample(nn.Module):
    def __init__(self, in_channels, with_conv, conv_op=ops.Conv2d, scale_factor=2.0):
@@ -2,196 +2,6 @@ import torch
 import math

 from .model import QwenImageTransformer2DModel
-from .model import QwenImageTransformerBlock
-
-
-class QwenImageFunControlBlock(QwenImageTransformerBlock):
-    def __init__(self, dim, num_attention_heads, attention_head_dim, has_before_proj=False, dtype=None, device=None, operations=None):
-        super().__init__(
-            dim=dim,
-            num_attention_heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-        self.has_before_proj = has_before_proj
-        if has_before_proj:
-            self.before_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
-        self.after_proj = operations.Linear(dim, dim, device=device, dtype=dtype)
-
-
-class QwenImageFunControlNetModel(torch.nn.Module):
-    def __init__(
-        self,
-        control_in_features=132,
-        inner_dim=3072,
-        num_attention_heads=24,
-        attention_head_dim=128,
-        num_control_blocks=5,
-        main_model_double=60,
-        injection_layers=(0, 12, 24, 36, 48),
-        dtype=None,
-        device=None,
-        operations=None,
-    ):
-        super().__init__()
-        self.dtype = dtype
-        self.main_model_double = main_model_double
-        self.injection_layers = tuple(injection_layers)
-        # Keep base hint scaling at 1.0 so user-facing strength behaves similarly
-        # to the reference Gen2/VideoX implementation around strength=1.
-        self.hint_scale = 1.0
-        self.control_img_in = operations.Linear(control_in_features, inner_dim, device=device, dtype=dtype)
-
-        self.control_blocks = torch.nn.ModuleList([])
-        for i in range(num_control_blocks):
-            self.control_blocks.append(
-                QwenImageFunControlBlock(
-                    dim=inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                    has_before_proj=(i == 0),
-                    dtype=dtype,
-                    device=device,
-                    operations=operations,
-                )
-            )
-
-    def _process_hint_tokens(self, hint):
-        if hint is None:
-            return None
-        if hint.ndim == 4:
-            hint = hint.unsqueeze(2)
-
-        # Fun checkpoints are trained with 33 latent channels before 2x2 packing:
-        # [control_latent(16), mask(1), inpaint_latent(16)] -> 132 features.
-        # Default behavior (no inpaint input in stock Apply ControlNet) should use
-        # zeros for mask/inpaint branches, matching VideoX fallback semantics.
-        expected_c = self.control_img_in.weight.shape[1] // 4
-        if hint.shape[1] == 16 and expected_c == 33:
-            zeros_mask = torch.zeros_like(hint[:, :1])
-            zeros_inpaint = torch.zeros_like(hint)
-            hint = torch.cat([hint, zeros_mask, zeros_inpaint], dim=1)
-
-        bs, c, t, h, w = hint.shape
-        hidden_states = torch.nn.functional.pad(hint, (0, w % 2, 0, h % 2))
-        orig_shape = hidden_states.shape
-        hidden_states = hidden_states.view(
-            orig_shape[0],
-            orig_shape[1],
-            orig_shape[-3],
-            orig_shape[-2] // 2,
-            2,
-            orig_shape[-1] // 2,
-            2,
-        )
-        hidden_states = hidden_states.permute(0, 2, 3, 5, 1, 4, 6)
-        hidden_states = hidden_states.reshape(
-            bs,
-            t * ((h + 1) // 2) * ((w + 1) // 2),
-            c * 4,
-        )
-
-        expected_in = self.control_img_in.weight.shape[1]
-        cur_in = hidden_states.shape[-1]
-        if cur_in < expected_in:
-            pad = torch.zeros(
-                (hidden_states.shape[0], hidden_states.shape[1], expected_in - cur_in),
-                device=hidden_states.device,
-                dtype=hidden_states.dtype,
-            )
-            hidden_states = torch.cat([hidden_states, pad], dim=-1)
-        elif cur_in > expected_in:
-            hidden_states = hidden_states[:, :, :expected_in]
-
-        return hidden_states
-
-    def forward(
-        self,
-        x,
-        timesteps,
-        context,
-        attention_mask=None,
-        guidance: torch.Tensor = None,
-        hint=None,
-        transformer_options={},
-        base_model=None,
-        **kwargs,
-    ):
-        if base_model is None:
-            raise RuntimeError("Qwen Fun ControlNet requires a QwenImage base model at runtime.")
-
-        encoder_hidden_states_mask = attention_mask
-        # Keep attention mask disabled inside Fun control blocks to mirror
-        # VideoX behavior (they rely on seq lengths for RoPE, not masked attention).
-        encoder_hidden_states_mask = None
-
-        hidden_states, img_ids, _ = base_model.process_img(x)
-        hint_tokens = self._process_hint_tokens(hint)
-        if hint_tokens is None:
-            raise RuntimeError("Qwen Fun ControlNet requires a control hint image.")
-
-        if hint_tokens.shape[1] != hidden_states.shape[1]:
-            max_tokens = min(hint_tokens.shape[1], hidden_states.shape[1])
-            hint_tokens = hint_tokens[:, :max_tokens]
-            hidden_states = hidden_states[:, :max_tokens]
-            img_ids = img_ids[:, :max_tokens]
-
-        txt_start = round(
-            max(
-                ((x.shape[-1] + (base_model.patch_size // 2)) // base_model.patch_size) // 2,
-                ((x.shape[-2] + (base_model.patch_size // 2)) // base_model.patch_size) // 2,
-            )
-        )
-        txt_ids = torch.arange(txt_start, txt_start + context.shape[1], device=x.device).reshape(1, -1, 1).repeat(x.shape[0], 1, 3)
-        ids = torch.cat((txt_ids, img_ids), dim=1)
-        image_rotary_emb = base_model.pe_embedder(ids).to(x.dtype).contiguous()
-
-        hidden_states = base_model.img_in(hidden_states)
-        encoder_hidden_states = base_model.txt_norm(context)
-        encoder_hidden_states = base_model.txt_in(encoder_hidden_states)
-
-        if guidance is not None:
-            guidance = guidance * 1000
-
-        temb = (
-            base_model.time_text_embed(timesteps, hidden_states)
-            if guidance is None
-            else base_model.time_text_embed(timesteps, guidance, hidden_states)
-        )
-
-        c = self.control_img_in(hint_tokens)
-
-        for i, block in enumerate(self.control_blocks):
-            if i == 0:
-                c_in = block.before_proj(c) + hidden_states
-                all_c = []
-            else:
-                all_c = list(torch.unbind(c, dim=0))
-                c_in = all_c.pop(-1)
-
-            encoder_hidden_states, c_out = block(
-                hidden_states=c_in,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_hidden_states_mask=encoder_hidden_states_mask,
-                temb=temb,
-                image_rotary_emb=image_rotary_emb,
-                transformer_options=transformer_options,
-            )
-
-            c_skip = block.after_proj(c_out) * self.hint_scale
-            all_c += [c_skip, c_out]
-            c = torch.stack(all_c, dim=0)
-
-        hints = torch.unbind(c, dim=0)[:-1]
-
-        controlnet_block_samples = [None] * self.main_model_double
-        for local_idx, base_idx in enumerate(self.injection_layers):
-            if local_idx < len(hints) and base_idx < len(controlnet_block_samples):
-                controlnet_block_samples[base_idx] = hints[local_idx]
-
-        return {"input": controlnet_block_samples}


 class QwenImageControlNetModel(QwenImageTransformer2DModel):
@@ -332,12 +332,6 @@ def model_lora_keys_unet(model, key_map={}):
                key_map["{}".format(key_lora)] = k
                key_map["transformer.{}".format(key_lora)] = k

-    if isinstance(model, comfy.model_base.ACEStep15):
-        for k in sdk:
-            if k.startswith("diffusion_model.decoder.") and k.endswith(".weight"):
-                key_lora = k[len("diffusion_model.decoder."):-len(".weight")]
-                key_map["base_model.model.{}".format(key_lora)] = k  # Official base model loras
-
    return key_map


@@ -374,31 +368,6 @@ def pad_tensor_to_shape(tensor: torch.Tensor, new_shape: list[int]) -> torch.Ten

    return padded_tensor

-def calculate_shape(patches, weight, key, original_weights=None):
-    current_shape = weight.shape
-
-    for p in patches:
-        v = p[1]
-        offset = p[3]
-
-        # Offsets restore the old shape; lists force a diff without metadata
-        if offset is not None or isinstance(v, list):
-            continue
-
-        if isinstance(v, weight_adapter.WeightAdapterBase):
-            adapter_shape = v.calculate_shape(key)
-            if adapter_shape is not None:
-                current_shape = adapter_shape
-            continue
-
-        # Standard diff logic with padding
-        if len(v) == 2:
-            patch_type, patch_data = v[0], v[1]
-            if patch_type == "diff" and len(patch_data) > 1 and patch_data[1]['pad_weight']:
-                current_shape = patch_data[0].shape
-
-    return current_shape
-
 def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, original_weights=None):
    for p in patches:
        strength = p[0]
@@ -5,7 +5,7 @@ import comfy.utils
 def convert_lora_bfl_control(sd): #BFL loras for Flux
    sd_out = {}
    for k in sd:
-        k_to = "diffusion_model.{}".format(k.replace(".lora_B.bias", ".diff_b").replace("_norm.scale", "_norm.set_weight"))
+        k_to = "diffusion_model.{}".format(k.replace(".lora_B.bias", ".diff_b").replace("_norm.scale", "_norm.scale.set_weight"))
        sd_out[k_to] = sd[k]

    sd_out["diffusion_model.img_in.reshape_weight"] = torch.tensor([sd["img_in.lora_B.weight"].shape[0], sd["img_in.lora_A.weight"].shape[1]])
@@ -1,81 +0,0 @@
-import math
-import torch
-from typing import NamedTuple
-
-from comfy.quant_ops import QuantizedTensor
-
-class TensorGeometry(NamedTuple):
-    shape: any
-    dtype: torch.dtype
-
-    def element_size(self):
-        info = torch.finfo(self.dtype) if self.dtype.is_floating_point else torch.iinfo(self.dtype)
-        return info.bits // 8
-
-    def numel(self):
-        return math.prod(self.shape)
-
-def tensors_to_geometries(tensors, dtype=None):
-    geometries = []
-    for t in tensors:
-        if t is None or isinstance(t, QuantizedTensor):
-            geometries.append(t)
-            continue
-        tdtype = t.dtype
-        if hasattr(t, "_model_dtype"):
-            tdtype = t._model_dtype
-        if dtype is not None:
-            tdtype = dtype
-        geometries.append(TensorGeometry(shape=t.shape, dtype=tdtype))
-    return geometries
-
-def vram_aligned_size(tensor):
-    if isinstance(tensor, list):
-        return sum([vram_aligned_size(t) for t in tensor])
-
-    if isinstance(tensor, QuantizedTensor):
-        inner_tensors, _ = tensor.__tensor_flatten__()
-        return vram_aligned_size([ getattr(tensor, attr) for attr in inner_tensors ])
-
-    if tensor is None:
-        return 0
-
-    size = tensor.numel() * tensor.element_size()
-    aligment_req = 1024
-    return (size + aligment_req - 1) // aligment_req * aligment_req
-
-def interpret_gathered_like(tensors, gathered):
-    offset = 0
-    dest_views = []
-
-    if gathered.dim() != 1 or gathered.element_size() != 1:
-        raise ValueError(f"Buffer must be 1D and single-byte (got {gathered.dim()}D {gathered.dtype})")
-
-    for tensor in tensors:
-
-        if tensor is None:
-            dest_views.append(None)
-            continue
-
-        if isinstance(tensor, QuantizedTensor):
-            inner_tensors, qt_ctx = tensor.__tensor_flatten__()
-            templates = { attr: getattr(tensor, attr) for attr in inner_tensors }
-        else:
-            templates = { "data": tensor }
-
-        actuals = {}
-        for attr, template in templates.items():
-            size = template.numel() * template.element_size()
-            if offset + size > gathered.numel():
-                raise ValueError(f"Buffer too small: needs {offset + size} bytes, but only has {gathered.numel()}. ")
-            actuals[attr] = gathered[offset:offset+size].view(dtype=template.dtype).view(template.shape)
-            offset += vram_aligned_size(template)
-
-        if isinstance(tensor, QuantizedTensor):
-            dest_views.append(QuantizedTensor.__tensor_unflatten__(actuals, qt_ctx, 0, 0))
-        else:
-            dest_views.append(actuals["data"])
-
-    return dest_views
-
-aimdo_allocator = None
@@ -50,7 +50,6 @@ import comfy.ldm.omnigen.omnigen2
 import comfy.ldm.qwen_image.model
 import comfy.ldm.kandinsky5.model
 import comfy.ldm.anima.model
-import comfy.ldm.ace.ace_step15

 import comfy.model_management
 import comfy.patcher_extension
@@ -147,8 +146,6 @@ class BaseModel(torch.nn.Module):
                self.diffusion_model.to(memory_format=torch.channels_last)
                logging.debug("using channels last mode for diffusion model")
            logging.info("model weight dtype {}, manual cast: {}".format(self.get_dtype(), self.manual_cast_dtype))
-            comfy.model_management.archive_model_dtypes(self.diffusion_model)
-
        self.model_type = model_type
        self.model_sampling = model_sampling(model_config, model_type)

@@ -178,7 +175,10 @@ class BaseModel(torch.nn.Module):
            xc = torch.cat([xc] + [comfy.model_management.cast_to_device(c_concat, xc.device, xc.dtype)], dim=1)

        context = c_crossattn
-        dtype = self.get_dtype_inference()
+        dtype = self.get_dtype()
+
+        if self.manual_cast_dtype is not None:
+            dtype = self.manual_cast_dtype

        xc = xc.to(dtype)
        device = xc.device
@@ -215,13 +215,6 @@ class BaseModel(torch.nn.Module):
    def get_dtype(self):
        return self.diffusion_model.dtype

-    def get_dtype_inference(self):
-        dtype = self.get_dtype()
-
-        if self.manual_cast_dtype is not None:
-            dtype = self.manual_cast_dtype
-        return dtype
-
    def encode_adm(self, **kwargs):
        return None

@@ -306,7 +299,7 @@ class BaseModel(torch.nn.Module):

        return out

-    def load_model_weights(self, sd, unet_prefix="", assign=False):
+    def load_model_weights(self, sd, unet_prefix=""):
        to_load = {}
        keys = list(sd.keys())
        for k in keys:
@@ -314,7 +307,7 @@ class BaseModel(torch.nn.Module):
                to_load[k[len(unet_prefix):]] = sd.pop(k)

        to_load = self.model_config.process_unet_state_dict(to_load)
-        m, u = self.diffusion_model.load_state_dict(to_load, strict=False, assign=assign)
+        m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
        if len(m) > 0:
            logging.warning("unet missing: {}".format(m))

@@ -329,7 +322,7 @@ class BaseModel(torch.nn.Module):
    def process_latent_out(self, latent):
        return self.latent_format.process_out(latent)

-    def state_dict_for_saving(self, unet_state_dict, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
+    def state_dict_for_saving(self, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
        extra_sds = []
        if clip_state_dict is not None:
            extra_sds.append(self.model_config.process_clip_state_dict_for_saving(clip_state_dict))
@@ -337,7 +330,10 @@ class BaseModel(torch.nn.Module):
            extra_sds.append(self.model_config.process_vae_state_dict_for_saving(vae_state_dict))
        if clip_vision_state_dict is not None:
            extra_sds.append(self.model_config.process_clip_vision_state_dict_for_saving(clip_vision_state_dict))
+
+        unet_state_dict = self.diffusion_model.state_dict()
        unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
+
        if self.model_type == ModelType.V_PREDICTION:
            unet_state_dict["v_pred"] = torch.tensor([])

@@ -376,7 +372,9 @@ class BaseModel(torch.nn.Module):
                    input_shapes += shape

        if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
-            dtype = self.get_dtype_inference()
+            dtype = self.get_dtype()
+            if self.manual_cast_dtype is not None:
+                dtype = self.manual_cast_dtype
            #TODO: this needs to be tweaked
            area = sum(map(lambda input_shape: input_shape[0] * math.prod(input_shape[2:]), input_shapes))
            return (area * comfy.model_management.dtype_size(dtype) * 0.01 * self.memory_usage_factor) * (1024 * 1024)
@@ -778,8 +776,8 @@ class StableAudio1(BaseModel):
            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
        return out

-    def state_dict_for_saving(self, unet_state_dict, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
-        sd = super().state_dict_for_saving(unet_state_dict, clip_state_dict=clip_state_dict, vae_state_dict=vae_state_dict, clip_vision_state_dict=clip_vision_state_dict)
+    def state_dict_for_saving(self, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
+        sd = super().state_dict_for_saving(clip_state_dict=clip_state_dict, vae_state_dict=vae_state_dict, clip_vision_state_dict=clip_vision_state_dict)
        d = {"conditioner.conditioners.seconds_start.": self.seconds_start_embedder.state_dict(), "conditioner.conditioners.seconds_total.": self.seconds_total_embedder.state_dict()}
        for k in d:
            s = d[k]
@@ -1162,16 +1160,12 @@ class Anima(BaseModel):
        device = kwargs["device"]
        if cross_attn is not None:
            if t5xxl_ids is not None:
+                cross_attn = self.diffusion_model.preprocess_text_embeds(cross_attn.to(device=device, dtype=self.get_dtype()), t5xxl_ids.unsqueeze(0).to(device=device))
                if t5xxl_weights is not None:
-                    t5xxl_weights = t5xxl_weights.unsqueeze(0).unsqueeze(-1).to(cross_attn)
-                t5xxl_ids = t5xxl_ids.unsqueeze(0)
-
-                if torch.is_inference_mode_enabled():  # if not we are training
-                    cross_attn = self.diffusion_model.preprocess_text_embeds(cross_attn.to(device=device, dtype=self.get_dtype_inference()), t5xxl_ids.to(device=device), t5xxl_weights=t5xxl_weights.to(device=device, dtype=self.get_dtype_inference()))
-                else:
-                    out['t5xxl_ids'] = comfy.conds.CONDRegular(t5xxl_ids)
-                    out['t5xxl_weights'] = comfy.conds.CONDRegular(t5xxl_weights)
+                    cross_attn *= t5xxl_weights.unsqueeze(0).unsqueeze(-1).to(cross_attn)

+                if cross_attn.shape[1] < 512:
+                    cross_attn = torch.nn.functional.pad(cross_attn, (0, 0, 0, 512 - cross_attn.shape[1]))
            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
        return out

@@ -1547,49 +1541,6 @@ class ACEStep(BaseModel):
        out['lyrics_strength'] = comfy.conds.CONDConstant(kwargs.get("lyrics_strength", 1.0))
        return out

-class ACEStep15(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ace.ace_step15.AceStepConditionGenerationModel)
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        device = kwargs["device"]
-        noise = kwargs["noise"]
-
-        cross_attn = kwargs.get("cross_attn", None)
-        if cross_attn is not None:
-            if torch.count_nonzero(cross_attn) == 0:
-                out['replace_with_null_embeds'] = comfy.conds.CONDConstant(True)
-            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
-
-        conditioning_lyrics = kwargs.get("conditioning_lyrics", None)
-        if cross_attn is not None:
-            out['lyric_embed'] = comfy.conds.CONDRegular(conditioning_lyrics)
-
-        refer_audio = kwargs.get("reference_audio_timbre_latents", None)
-        if refer_audio is None or len(refer_audio) == 0:
-            refer_audio = comfy.ldm.ace.ace_step15.get_silence_latent(noise.shape[2], device)
-            pass_audio_codes = True
-        else:
-            refer_audio = refer_audio[-1][:, :, :noise.shape[2]]
-            out['is_covers'] = comfy.conds.CONDConstant(True)
-            pass_audio_codes = False
-
-        if pass_audio_codes:
-            audio_codes = kwargs.get("audio_codes", None)
-            if audio_codes is not None:
-                out['audio_codes'] = comfy.conds.CONDRegular(torch.tensor(audio_codes, device=device))
-                refer_audio = refer_audio[:, :, :750]
-            else:
-                out['is_covers'] = comfy.conds.CONDConstant(False)
-
-        if refer_audio.shape[2] < noise.shape[2]:
-            pad = comfy.ldm.ace.ace_step15.get_silence_latent(noise.shape[2], device)
-            refer_audio = torch.cat([refer_audio.to(pad), pad[:, :, refer_audio.shape[2]:]], dim=2)
-
-        out['refer_audio'] = comfy.conds.CONDRegular(refer_audio)
-        return out
-
 class Omnigen2(BaseModel):
    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.omnigen.omnigen2.OmniGen2Transformer2DModel)
@@ -19,12 +19,6 @@ def count_blocks(state_dict_keys, prefix_string):
        count += 1
    return count

-def any_suffix_in(keys, prefix, main, suffix_list=[]):
-    for x in suffix_list:
-        if "{}{}{}".format(prefix, main, x) in keys:
-            return True
-    return False
-
 def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
    context_dim = None
    use_linear_in_transformer = False
@@ -192,7 +186,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
            dit_config["meanflow_sum"] = False
        return dit_config

-    if any_suffix_in(state_dict_keys, key_prefix, 'double_blocks.0.img_attn.norm.key_norm.', ["weight", "scale"]) and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.norms.0.', ["weight", "scale"])): #Flux, Chroma or Chroma Radiance (has no img_in.weight)
+    if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or f"{key_prefix}distilled_guidance_layer.norms.0.scale" in state_dict_keys): #Flux, Chroma or Chroma Radiance (has no img_in.weight)
        dit_config = {}
        if '{}double_stream_modulation_img.lin.weight'.format(key_prefix) in state_dict_keys:
            dit_config["image_model"] = "flux2"
@@ -247,8 +241,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):

        dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
        dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')
-
-        if any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.0.norms.0.', ["weight", "scale"]) or any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.norms.0.', ["weight", "scale"]): #Chroma
+        if '{}distilled_guidance_layer.0.norms.0.scale'.format(key_prefix) in state_dict_keys or '{}distilled_guidance_layer.norms.0.scale'.format(key_prefix) in state_dict_keys: #Chroma
            dit_config["image_model"] = "chroma"
            dit_config["in_channels"] = 64
            dit_config["out_channels"] = 64
@@ -256,8 +249,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
            dit_config["out_dim"] = 3072
            dit_config["hidden_dim"] = 5120
            dit_config["n_layers"] = 5
-
-            if any_suffix_in(state_dict_keys, key_prefix, 'nerf_blocks.0.norm.', ["weight", "scale"]): #Chroma Radiance
+            if f"{key_prefix}nerf_blocks.0.norm.scale" in state_dict_keys: #Chroma Radiance
                dit_config["image_model"] = "chroma_radiance"
                dit_config["in_channels"] = 3
                dit_config["out_channels"] = 3
@@ -267,7 +259,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                dit_config["nerf_depth"] = 4
                dit_config["nerf_max_freqs"] = 8
                dit_config["nerf_tile_size"] = 512
-                dit_config["nerf_final_head_type"] = "conv" if any_suffix_in(state_dict_keys, key_prefix, 'nerf_final_layer_conv.norm.', ["weight", "scale"]) else "linear"
+                dit_config["nerf_final_head_type"] = "conv" if f"{key_prefix}nerf_final_layer_conv.norm.scale" in state_dict_keys else "linear"
                dit_config["nerf_embedder_dtype"] = torch.float32
                if "{}__x0__".format(key_prefix) in state_dict_keys: # x0 pred
                    dit_config["use_x0"] = True
@@ -276,7 +268,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
        else:
            dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
            dit_config["yak_mlp"] = '{}double_blocks.0.img_mlp.gate_proj.weight'.format(key_prefix) in state_dict_keys
-            dit_config["txt_norm"] = any_suffix_in(state_dict_keys, key_prefix, 'txt_norm.', ["weight", "scale"])
+            dit_config["txt_norm"] = "{}txt_norm.scale".format(key_prefix) in state_dict_keys
            if dit_config["yak_mlp"] and dit_config["txt_norm"]:  # Ovis model
                dit_config["txt_ids_dims"] = [1, 2]

@@ -663,11 +655,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
        dit_config["num_visual_blocks"] = count_blocks(state_dict_keys, '{}visual_transformer_blocks.'.format(key_prefix) + '{}.')
        return dit_config

-    if '{}encoder.lyric_encoder.layers.0.input_layernorm.weight'.format(key_prefix) in state_dict_keys:
-        dit_config = {}
-        dit_config["audio_model"] = "ace1.5"
-        return dit_config
-
    if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
        return None

@@ -20,20 +20,12 @@ import psutil
 import logging
 from enum import Enum
 from comfy.cli_args import args, PerformanceFeature
-import threading
 import torch
 import sys
 import platform
 import weakref
 import gc
 import os
-from contextlib import nullcontext
-import comfy.memory_management
-import comfy.utils
-import comfy.quant_ops
-
-import comfy_aimdo.torch
-import comfy_aimdo.model_vbar

 class VRAMState(Enum):
    DISABLED = 0    #No vram present: no need to move models to vram
@@ -55,11 +47,6 @@ cpu_state = CPUState.GPU

 total_vram = 0

-
-# Training Related State
-in_training = False
-
-
 def get_supported_float8_types():
    float8_types = []
    try:
@@ -591,15 +578,9 @@ WINDOWS = any(platform.win32_ver())

 EXTRA_RESERVED_VRAM = 400 * 1024 * 1024
 if WINDOWS:
-    import comfy.windows
    EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue
    if total_vram > (15 * 1024):  # more extra reserved vram on 16GB+ cards
        EXTRA_RESERVED_VRAM += 100 * 1024 * 1024
-    def get_free_ram():
-        return comfy.windows.get_free_ram()
-else:
-    def get_free_ram():
-        return psutil.virtual_memory().available

 if args.reserve_vram is not None:
    EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024
@@ -611,7 +592,7 @@ def extra_reserved_memory():
 def minimum_inference_memory():
    return (1024 * 1024 * 1024) * 0.8 + extra_reserved_memory()

-def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, ram_required=0):
+def free_memory(memory_required, device, keep_loaded=[]):
    cleanup_models_gc()
    unloaded_model = []
    can_unload = []
@@ -626,23 +607,15 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, ram_

    for x in sorted(can_unload):
        i = x[-1]
-        memory_to_free = 1e32
-        ram_to_free = 1e32
+        memory_to_free = None
        if not DISABLE_SMART_MEMORY:
-            memory_to_free = memory_required - get_free_memory(device)
-            ram_to_free = ram_required - get_free_ram()
-
-        if current_loaded_models[i].model.is_dynamic() and for_dynamic:
-            #don't actually unload dynamic models for the sake of other dynamic models
-            #as that works on-demand.
-            memory_required -= current_loaded_models[i].model.loaded_size()
-            memory_to_free = 0
-        if memory_to_free > 0 and current_loaded_models[i].model_unload(memory_to_free):
-            logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
+            free_mem = get_free_memory(device)
+            if free_mem > memory_required:
+                break
+            memory_to_free = memory_required - free_mem
+        logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
+        if current_loaded_models[i].model_unload(memory_to_free):
            unloaded_model.append(i)
-        if ram_to_free > 0:
-            logging.debug(f"RAM Unloading {current_loaded_models[i].model.model.__class__.__name__}")
-            current_loaded_models[i].model.partially_unload_ram(ram_to_free)

    for i in sorted(unloaded_model, reverse=True):
        unloaded_models.append(current_loaded_models.pop(i))
@@ -677,10 +650,7 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu

    models_to_load = []

-    free_for_dynamic=True
    for x in models:
-        if not x.is_dynamic():
-            free_for_dynamic = False
        loaded_model = LoadedModel(x)
        try:
            loaded_model_index = current_loaded_models.index(loaded_model)
@@ -706,25 +676,19 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
            model_to_unload.model.detach(unpatch_all=False)
            model_to_unload.model_finalizer.detach()

-
    total_memory_required = {}
-    total_ram_required = {}
    for loaded_model in models_to_load:
        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
-        #x2, one to make sure the OS can fit the model for loading in disk cache, and for us to do any pinning we
-        #want to do.
-        #FIXME: This should subtract off the to_load current pin consumption.
-        total_ram_required[loaded_model.device] = total_ram_required.get(loaded_model.device, 0) + loaded_model.model_memory() * 2

    for device in total_memory_required:
        if device != torch.device("cpu"):
-            free_memory(total_memory_required[device] * 1.1 + extra_mem, device, for_dynamic=free_for_dynamic, ram_required=total_ram_required[device])
+            free_memory(total_memory_required[device] * 1.1 + extra_mem, device)

    for device in total_memory_required:
        if device != torch.device("cpu"):
            free_mem = get_free_memory(device)
            if free_mem < minimum_memory_required:
-                models_l = free_memory(minimum_memory_required, device, for_dynamic=free_for_dynamic)
+                models_l = free_memory(minimum_memory_required, device)
                logging.info("{} models unloaded.".format(len(models_l)))

    for loaded_model in models_to_load:
@@ -768,9 +732,6 @@ def loaded_models(only_currently_used=False):

 def cleanup_models_gc():
    do_gc = False
-
-    reset_cast_buffers()
-
    for i in range(len(current_loaded_models)):
        cur = current_loaded_models[i]
        if cur.is_dead():
@@ -788,11 +749,6 @@ def cleanup_models_gc():
                logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__))


-def archive_model_dtypes(model):
-    for name, module in model.named_modules():
-        for param_name, param in module.named_parameters(recurse=False):
-            setattr(module, f"{param_name}_comfy_model_dtype", param.dtype)
-

 def cleanup_models():
    to_delete = []
@@ -836,7 +792,7 @@ def unet_inital_load_device(parameters, dtype):

    mem_dev = get_free_memory(torch_dev)
    mem_cpu = get_free_memory(cpu_dev)
-    if mem_dev > mem_cpu and model_size < mem_dev and comfy.memory_management.aimdo_allocator is None:
+    if mem_dev > mem_cpu and model_size < mem_dev:
        return torch_dev
    else:
        return cpu_dev
@@ -1095,51 +1051,6 @@ def current_stream(device):
        return None

 stream_counters = {}
-
-STREAM_CAST_BUFFERS = {}
-LARGEST_CASTED_WEIGHT = (None, 0)
-
-def get_cast_buffer(offload_stream, device, size, ref):
-    global LARGEST_CASTED_WEIGHT
-
-    if offload_stream is not None:
-        wf_context = offload_stream
-        if hasattr(wf_context, "as_context"):
-            wf_context = wf_context.as_context(offload_stream)
-    else:
-        wf_context = nullcontext()
-
-    cast_buffer = STREAM_CAST_BUFFERS.get(offload_stream, None)
-    if cast_buffer is None or cast_buffer.numel() < size:
-        if ref is LARGEST_CASTED_WEIGHT[0]:
-            #If there is one giant weight we do not want both streams to
-            #allocate a buffer for it. It's up to the caster to get the other
-            #offload stream in this corner case
-            return None
-        if cast_buffer is not None and cast_buffer.numel() > 50 * (1024 ** 2):
-            #I want my wrongly sized 50MB+ of VRAM back from the caching allocator right now
-            synchronize()
-            del STREAM_CAST_BUFFERS[offload_stream]
-            del cast_buffer
-            #FIXME: This doesn't work in Aimdo because mempool cant clear cache
-            soft_empty_cache()
-        with wf_context:
-            cast_buffer = torch.empty((size), dtype=torch.int8, device=device)
-            STREAM_CAST_BUFFERS[offload_stream] = cast_buffer
-
-        if  size > LARGEST_CASTED_WEIGHT[1]:
-            LARGEST_CASTED_WEIGHT = (ref, size)
-
-    return cast_buffer
-
-def reset_cast_buffers():
-    global LARGEST_CASTED_WEIGHT
-    LARGEST_CASTED_WEIGHT = (None, 0)
-    for offload_stream in STREAM_CAST_BUFFERS:
-        offload_stream.synchronize()
-    STREAM_CAST_BUFFERS.clear()
-    soft_empty_cache()
-
 def get_offload_stream(device):
    stream_counter = stream_counters.get(device, 0)
    if NUM_STREAMS == 0:
@@ -1182,61 +1093,7 @@ def sync_stream(device, stream):
        return
    current_stream(device).wait_stream(stream)

-
-def cast_to_gathered(tensors, r, non_blocking=False, stream=None):
-    wf_context = nullcontext()
-    if stream is not None:
-       wf_context = stream
-       if hasattr(wf_context, "as_context"):
-           wf_context = wf_context.as_context(stream)
-
-    dest_views = comfy.memory_management.interpret_gathered_like(tensors, r)
-    with wf_context:
-        for tensor in tensors:
-            dest_view = dest_views.pop(0)
-            if tensor is None:
-                continue
-            dest_view.copy_(tensor, non_blocking=non_blocking)
-
-
-def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None, r=None):
-    if hasattr(weight, "_v"):
-        #Unexpected usage patterns. There is no reason these don't work but they
-        #have no testing and no callers do this.
-        assert r is None
-        assert stream is None
-
-        cast_geometry = comfy.memory_management.tensors_to_geometries([ weight ])
-
-        if dtype is None:
-            dtype = weight._model_dtype
-
-        signature = comfy_aimdo.model_vbar.vbar_fault(weight._v)
-        if signature is not None:
-            if comfy_aimdo.model_vbar.vbar_signature_compare(signature, weight._v_signature):
-                v_tensor = weight._v_tensor
-            else:
-                raw_tensor = comfy_aimdo.torch.aimdo_to_tensor(weight._v, device)
-                v_tensor = comfy.memory_management.interpret_gathered_like(cast_geometry, raw_tensor)[0]
-                weight._v_tensor = v_tensor
-                weight._v_signature = signature
-                #Send it over
-                v_tensor.copy_(weight, non_blocking=non_blocking)
-            return v_tensor.to(dtype=dtype)
-
-        r = torch.empty_like(weight, dtype=dtype, device=device)
-
-        if weight.dtype != r.dtype and weight.dtype != weight._model_dtype:
-            #Offloaded casting could skip this, however it would make the quantizations
-            #inconsistent between loaded and offloaded weights. So force the double casting
-            #that would happen in regular flow to make offload deterministic.
-            cast_buffer = torch.empty_like(weight, dtype=weight._model_dtype, device=device)
-            cast_buffer.copy_(weight, non_blocking=non_blocking)
-            weight = cast_buffer
-        r.copy_(weight, non_blocking=non_blocking)
-
-        return r
-
+def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None):
    if device is None or weight.device == device:
        if not copy:
            if dtype is None or weight.dtype == dtype:
@@ -1255,12 +1112,10 @@ def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, str
        if hasattr(wf_context, "as_context"):
            wf_context = wf_context.as_context(stream)
        with wf_context:
-            if r is None:
-                r = torch.empty_like(weight, dtype=dtype, device=device)
+            r = torch.empty_like(weight, dtype=dtype, device=device)
            r.copy_(weight, non_blocking=non_blocking)
    else:
-        if r is None:
-            r = torch.empty_like(weight, dtype=dtype, device=device)
+        r = torch.empty_like(weight, dtype=dtype, device=device)
        r.copy_(weight, non_blocking=non_blocking)
    return r

@@ -1280,14 +1135,14 @@ if not args.disable_pinned_memory:
            MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95
        logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))

-PINNING_ALLOWED_TYPES = set(["Tensor", "Parameter", "QuantizedTensor"])
+PINNING_ALLOWED_TYPES = set(["Parameter", "QuantizedTensor"])

 def discard_cuda_async_error():
    try:
        a = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
        b = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
        _ = a + b
-        synchronize()
+        torch.cuda.synchronize()
    except torch.AcceleratorError:
        #Dump it! We already know about it from the synchronous return
        pass
@@ -1691,12 +1546,6 @@ def lora_compute_dtype(device):
    LORA_COMPUTE_DTYPES[device] = dtype
    return dtype

-def synchronize():
-    if is_intel_xpu():
-        torch.xpu.synchronize()
-    elif torch.cuda.is_available():
-        torch.cuda.synchronize()
-
 def soft_empty_cache(force=False):
    global cpu_state
    if cpu_state == CPUState.MPS:
@@ -1708,7 +1557,6 @@ def soft_empty_cache(force=False):
    elif is_mlu():
        torch.mlu.empty_cache()
    elif torch.cuda.is_available():
-        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

@@ -1720,6 +1568,9 @@ def debug_memory_summary():
        return torch.cuda.memory.memory_summary()
    return ""

+#TODO: might be cleaner to put this somewhere else
+import threading
+
 class InterruptProcessingException(Exception):
    pass

@@ -19,6 +19,7 @@
 from __future__ import annotations

 import collections
+import copy
 import inspect
 import logging
 import math
@@ -37,7 +38,19 @@ from comfy.comfy_types import UnetWrapperFunction
 from comfy.quant_ops import QuantizedTensor
 from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP

-import comfy_aimdo.model_vbar
+
+def string_to_seed(data):
+    crc = 0xFFFFFFFF
+    for byte in data:
+        if isinstance(byte, str):
+            byte = ord(byte)
+        crc ^= byte
+        for _ in range(8):
+            if crc & 1:
+                crc = (crc >> 1) ^ 0xEDB88320
+            else:
+                crc >>= 1
+    return crc ^ 0xFFFFFFFF

 def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
    to = model_options["transformer_options"].copy()
@@ -110,10 +123,6 @@ def move_weight_functions(m, device):
                memory += f.move_to(device=device)
    return memory

-def string_to_seed(data):
-    logging.warning("WARNING: string_to_seed has moved from comfy.model_patcher to comfy.utils")
-    return comfy.utils.string_to_seed(data)
-
 class LowVramPatch:
    def __init__(self, key, patches, convert_func=None, set_func=None):
        self.key = key
@@ -160,11 +169,6 @@ def get_key_weight(model, key):

    return weight, set_func, convert_func

-def key_param_name_to_key(key, param):
-    if len(key) == 0:
-        return param
-    return "{}.{}".format(key, param)
-
 class AutoPatcherEjector:
    def __init__(self, model: 'ModelPatcher', skip_and_inject_on_exit_only=False):
        self.model = model
@@ -208,27 +212,6 @@ class MemoryCounter:
    def decrement(self, used: int):
        self.value -= used

-CustomTorchDevice = collections.namedtuple("FakeDevice", ["type", "index"])("comfy-lazy-caster", 0)
-
-class LazyCastingParam(torch.nn.Parameter):
-    def __new__(cls, model, key, tensor):
-        return super().__new__(cls, tensor)
-
-    def __init__(self, model, key, tensor):
-        self.model = model
-        self.key = key
-
-    @property
-    def device(self):
-        return CustomTorchDevice
-
-    #safetensors will .to() us to the cpu which we catch here to cast on demand. The returned tensor is
-    #then just a short lived thing in the safetensors serialization logic inside its big for loop over
-    #all weights getting garbage collected per-weight
-    def to(self, *args, **kwargs):
-        return self.model.patch_weight_to_device(self.key, device_to=self.model.load_device, return_weight=True).to("cpu")
-
-
 class ModelPatcher:
    def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
        self.size = size
@@ -286,9 +269,6 @@ class ModelPatcher:
        if not hasattr(self.model, 'model_offload_buffer_memory'):
            self.model.model_offload_buffer_memory = 0

-    def is_dynamic(self):
-        return False
-
    def model_size(self):
        if self.size > 0:
            return self.size
@@ -304,9 +284,6 @@ class ModelPatcher:
    def lowvram_patch_counter(self):
        return self.model.lowvram_patch_counter

-    def get_free_memory(self, device):
-        return comfy.model_management.get_free_memory(device)
-
    def clone(self):
        n = self.__class__(self.model, self.load_device, self.offload_device, self.model_size(), weight_inplace_update=self.weight_inplace_update)
        n.patches = {}
@@ -316,7 +293,7 @@ class ModelPatcher:

        n.object_patches = self.object_patches.copy()
        n.weight_wrapper_patches = self.weight_wrapper_patches.copy()
-        n.model_options = comfy.utils.deepcopy_list_dict(self.model_options)
+        n.model_options = copy.deepcopy(self.model_options)
        n.backup = self.backup
        n.object_patches_backup = self.object_patches_backup
        n.parent = self
@@ -406,16 +383,13 @@ class ModelPatcher:
    def memory_required(self, input_shape):
        return self.model.memory_required(input_shape=input_shape)

-    def disable_model_cfg1_optimization(self):
-        self.model_options["disable_cfg1_optimization"] = True
-
    def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):
        if len(inspect.signature(sampler_cfg_function).parameters) == 3:
            self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
        else:
            self.model_options["sampler_cfg_function"] = sampler_cfg_function
        if disable_cfg1_optimization:
-            self.disable_model_cfg1_optimization()
+            self.model_options["disable_cfg1_optimization"] = True

    def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
        self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)
@@ -637,14 +611,14 @@ class ModelPatcher:
                        sd.pop(k)
            return sd

-    def patch_weight_to_device(self, key, device_to=None, inplace_update=False, return_weight=False):
-        weight, set_func, convert_func = get_key_weight(self.model, key)
+    def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
        if key not in self.patches:
-            return weight
+            return

+        weight, set_func, convert_func = get_key_weight(self.model, key)
        inplace_update = self.weight_inplace_update or inplace_update

-        if key not in self.backup and not return_weight:
+        if key not in self.backup:
            self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)

        temp_dtype = comfy.model_management.lora_compute_dtype(device_to)
@@ -657,15 +631,13 @@ class ModelPatcher:

        out_weight = comfy.lora.calculate_weight(self.patches[key], temp_weight, key)
        if set_func is None:
-            out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=comfy.utils.string_to_seed(key))
-            if return_weight:
-                return out_weight
-            elif inplace_update:
+            out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))
+            if inplace_update:
                comfy.utils.copy_to_param(self.model, key, out_weight)
            else:
                comfy.utils.set_attr_param(self.model, key, out_weight)
        else:
-            return set_func(out_weight, inplace_update=inplace_update, seed=comfy.utils.string_to_seed(key), return_weight=return_weight)
+            set_func(out_weight, inplace_update=inplace_update, seed=string_to_seed(key))

    def pin_weight_to_device(self, key):
        weight, set_func, convert_func = get_key_weight(self.model, key)
@@ -682,19 +654,18 @@ class ModelPatcher:
        for key in list(self.pinned):
            self.unpin_weight(key)

-    def _load_list(self, prio_comfy_cast_weights=False, default_device=None):
+    def _load_list(self):
        loading = []
        for n, m in self.model.named_modules():
-            default = False
-            params = { name: param for name, param in m.named_parameters(recurse=False) }
+            params = []
+            skip = False
+            for name, param in m.named_parameters(recurse=False):
+                params.append(name)
            for name, param in m.named_parameters(recurse=True):
                if name not in params:
-                    default = True # default random weights in non leaf modules
+                    skip = True # skip random weights in non leaf modules
                    break
-            if default and default_device is not None:
-                for param in params.values():
-                    param.data = param.data.to(device=default_device)
-            if not default and (hasattr(m, "comfy_cast_weights") or len(params) > 0):
+            if not skip and (hasattr(m, "comfy_cast_weights") or len(params) > 0):
                module_mem = comfy.model_management.module_size(m)
                module_offload_mem = module_mem
                if hasattr(m, "comfy_cast_weights"):
@@ -710,8 +681,7 @@ class ModelPatcher:
                        return 0
                    module_offload_mem += check_module_offload_mem("{}.weight".format(n))
                    module_offload_mem += check_module_offload_mem("{}.bias".format(n))
-                prepend = (not hasattr(m, "comfy_cast_weights"),) if prio_comfy_cast_weights else ()
-                loading.append(prepend + (module_offload_mem, module_mem, n, m, params))
+                loading.append((module_offload_mem, module_mem, n, m, params))
        return loading

    def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
@@ -803,7 +773,7 @@ class ModelPatcher:
                        continue

                for param in params:
-                    key = key_param_name_to_key(n, param)
+                    key = "{}.{}".format(n, param)
                    self.unpin_weight(key)
                    self.patch_weight_to_device(key, device_to=device_to)
                if comfy.model_management.is_device_cuda(device_to):
@@ -819,7 +789,7 @@ class ModelPatcher:
                n = x[1]
                params = x[3]
                for param in params:
-                    self.pin_weight_to_device(key_param_name_to_key(n, param))
+                    self.pin_weight_to_device("{}.{}".format(n, param))

            usable_stat = "{:.2f} MB usable,".format(lowvram_model_memory / (1024 * 1024)) if lowvram_model_memory < 1e32 else ""
            if lowvram_counter > 0:
@@ -925,7 +895,7 @@ class ModelPatcher:
                if hasattr(m, "comfy_patched_weights") and m.comfy_patched_weights == True:
                    move_weight = True
                    for param in params:
-                        key = key_param_name_to_key(n, param)
+                        key = "{}.{}".format(n, param)
                        bk = self.backup.get(key, None)
                        if bk is not None:
                            if not lowvram_possible:
@@ -976,7 +946,7 @@ class ModelPatcher:
                        logging.debug("freed {}".format(n))

                        for param in params:
-                            self.pin_weight_to_device(key_param_name_to_key(n, param))
+                            self.pin_weight_to_device("{}.{}".format(n, param))


            self.model.model_lowvram = True
@@ -1014,9 +984,6 @@ class ModelPatcher:

            return self.model.model_loaded_weight_memory - current_used

-    def partially_unload_ram(self, ram_to_unload):
-        pass
-
    def detach(self, unpatch_all=True):
        self.eject_model()
        self.model_patches_to(self.offload_device)
@@ -1350,10 +1317,10 @@ class ModelPatcher:
                                                 key, original_weights=original_weights)
        del original_weights[key]
        if set_func is None:
-            out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=comfy.utils.string_to_seed(key))
+            out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))
            comfy.utils.copy_to_param(self.model, key, out_weight)
        else:
-            set_func(out_weight, inplace_update=True, seed=comfy.utils.string_to_seed(key))
+            set_func(out_weight, inplace_update=True, seed=string_to_seed(key))
        if self.hook_mode == comfy.hooks.EnumHookMode.MaxSpeed:
            # TODO: disable caching if not enough system RAM to do so
            target_device = self.offload_device
@@ -1388,275 +1355,7 @@ class ModelPatcher:
        self.unpatch_hooks()
        self.clear_cached_hook_weights()

-    def state_dict_for_saving(self, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
-        unet_state_dict = self.model.diffusion_model.state_dict()
-        for k, v in unet_state_dict.items():
-            op_keys = k.rsplit('.', 1)
-            if (len(op_keys) < 2) or op_keys[1] not in ["weight", "bias"]:
-                continue
-            try:
-                op = comfy.utils.get_attr(self.model.diffusion_model, op_keys[0])
-            except:
-                continue
-            if not op or not hasattr(op, "comfy_cast_weights") or \
-                (hasattr(op, "comfy_patched_weights") and op.comfy_patched_weights == True):
-                continue
-            key = "diffusion_model." + k
-            unet_state_dict[k] = LazyCastingParam(self, key, comfy.utils.get_attr(self.model, key))
-        return self.model.state_dict_for_saving(unet_state_dict, clip_state_dict=clip_state_dict, vae_state_dict=vae_state_dict, clip_vision_state_dict=clip_vision_state_dict)
-
    def __del__(self):
        self.unpin_all_weights()
        self.detach(unpatch_all=False)

-class ModelPatcherDynamic(ModelPatcher):
-
-    def __new__(cls, model=None, load_device=None, offload_device=None, size=0, weight_inplace_update=False):
-        if load_device is not None and comfy.model_management.is_device_cpu(load_device):
-            #reroute to default MP for CPUs
-            return ModelPatcher(model, load_device, offload_device, size, weight_inplace_update)
-        return super().__new__(cls)
-
-    def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
-        super().__init__(model, load_device, offload_device, size, weight_inplace_update)
-        #this is now way more dynamic and we dont support the same base model for both Dynamic
-        #and non-dynamic patchers.
-        if hasattr(self.model, "model_loaded_weight_memory"):
-            del self.model.model_loaded_weight_memory
-        if not hasattr(self.model, "dynamic_vbars"):
-            self.model.dynamic_vbars = {}
-        assert load_device is not None
-
-    def is_dynamic(self):
-        return True
-
-    def _vbar_get(self, create=False):
-        if self.load_device == torch.device("cpu"):
-            return None
-        vbar = self.model.dynamic_vbars.get(self.load_device, None)
-        if create and vbar is None:
-            # x10. We dont know what model defined type casts we have in the vbar, but virtual address
-            # space is pretty free. This will cover someone casting an entire model from FP4 to FP32
-            # with some left over.
-            vbar = comfy_aimdo.model_vbar.ModelVBAR(self.model_size() * 10, self.load_device.index)
-            self.model.dynamic_vbars[self.load_device] = vbar
-        return vbar
-
-    def loaded_size(self):
-        vbar = self._vbar_get()
-        if vbar is None:
-            return 0
-        return vbar.loaded_size()
-
-    def get_free_memory(self, device):
-        #NOTE: on high condition / batch counts, estimate should have already vacated
-        #all non-dynamic models so this is safe even if its not 100% true that this
-        #would all be avaiable for inference use.
-        return comfy.model_management.get_total_memory(device) - self.model_size()
-
-    #Pinning is deferred to ops time. Assert against this API to avoid pin leaks.
-
-    def pin_weight_to_device(self, key):
-        raise RuntimeError("pin_weight_to_device invalid for dymamic weight loading")
-
-    def unpin_weight(self, key):
-        raise RuntimeError("unpin_weight invalid for dymamic weight loading")
-
-    def unpin_all_weights(self):
-        self.partially_unload_ram(1e32)
-
-    def memory_required(self, input_shape):
-        #Pad this significantly. We are trying to get away from precise estimates. This
-        #estimate is only used when using the ModelPatcherDynamic after ModelPatcher. If you
-        #use all ModelPatcherDynamic this is ignored and its all done dynamically.
-        return super().memory_required(input_shape=input_shape) * 1.3 + (1024 ** 3)
-
-
-    def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False, dirty=False):
-
-        #Force patching doesn't make sense in Dynamic loading, as you dont know what does and
-        #doesn't need to be forced at this stage. The only thing you could do would be patch
-        #it all on CPU which consumes huge RAM.
-        assert not force_patch_weights
-
-        #Full load doesn't make sense as we dont actually have any loader capability here and
-        #now.
-        assert not full_load
-
-        assert device_to == self.load_device
-
-        num_patches = 0
-        allocated_size = 0
-
-        with self.use_ejected():
-            self.unpatch_hooks()
-
-            vbar = self._vbar_get(create=True)
-            if vbar is not None:
-                vbar.prioritize()
-
-            #We force reserve VRAM for the non comfy-weight so we dont have to deal
-            #with pin and unpin syncrhonization which can be expensive for small weights
-            #with a high layer rate (e.g. autoregressive LLMs).
-            #prioritize the non-comfy weights (note the order reverse).
-            loading = self._load_list(prio_comfy_cast_weights=True, default_device=device_to)
-            loading.sort(reverse=True)
-
-            for x in loading:
-                _, _, _, n, m, params = x
-
-                def set_dirty(item, dirty):
-                    if dirty or not hasattr(item, "_v_signature"):
-                        item._v_signature = None
-
-                def setup_param(self, m, n, param_key):
-                    nonlocal num_patches
-                    key = key_param_name_to_key(n, param_key)
-
-                    weight_function = []
-
-                    weight, _, _ = get_key_weight(self.model, key)
-                    if weight is None:
-                        return (False, 0)
-                    if key in self.patches:
-                        if comfy.lora.calculate_shape(self.patches[key], weight, key) != weight.shape:
-                            return (True, 0)
-                        setattr(m, param_key + "_lowvram_function", LowVramPatch(key, self.patches))
-                        num_patches += 1
-                    else:
-                        setattr(m, param_key + "_lowvram_function", None)
-
-                    if key in self.weight_wrapper_patches:
-                        weight_function.extend(self.weight_wrapper_patches[key])
-                    setattr(m, param_key + "_function", weight_function)
-                    geometry = weight
-                    if not isinstance(weight, QuantizedTensor):
-                        model_dtype = getattr(m, param_key + "_comfy_model_dtype", None) or weight.dtype
-                        weight._model_dtype = model_dtype
-                        geometry = comfy.memory_management.TensorGeometry(shape=weight.shape, dtype=model_dtype)
-                    return (False, comfy.memory_management.vram_aligned_size(geometry))
-
-                def force_load_param(self, param_key, device_to):
-                    key = key_param_name_to_key(n, param_key)
-                    if key in self.backup:
-                        comfy.utils.set_attr_param(self.model, key, self.backup[key].weight)
-                    self.patch_weight_to_device(key, device_to=device_to)
-
-                if hasattr(m, "comfy_cast_weights"):
-                    m.comfy_cast_weights = True
-                    m.pin_failed = False
-                    m.seed_key = n
-                    set_dirty(m, dirty)
-
-                    force_load, v_weight_size = setup_param(self, m, n, "weight")
-                    force_load_bias, v_weight_bias = setup_param(self, m, n, "bias")
-                    force_load = force_load or force_load_bias
-                    v_weight_size += v_weight_bias
-
-                    if force_load:
-                        logging.info(f"Module {n} has resizing Lora - force loading")
-                        force_load_param(self, "weight", device_to)
-                        force_load_param(self, "bias", device_to)
-                    else:
-                        if vbar is not None and not hasattr(m, "_v"):
-                            m._v = vbar.alloc(v_weight_size)
-                        allocated_size += v_weight_size
-
-                else:
-                    for param in params:
-                        key = key_param_name_to_key(n, param)
-                        weight, _, _ = get_key_weight(self.model, key)
-                        weight.seed_key = key
-                        set_dirty(weight, dirty)
-                        geometry = weight
-                        model_dtype = getattr(m, param + "_comfy_model_dtype", None) or weight.dtype
-                        geometry = comfy.memory_management.TensorGeometry(shape=weight.shape, dtype=model_dtype)
-                        weight_size = geometry.numel() * geometry.element_size()
-                        if vbar is not None and not hasattr(weight, "_v"):
-                            weight._v = vbar.alloc(weight_size)
-                            weight._model_dtype = model_dtype
-                        allocated_size += weight_size
-                    vbar.set_watermark_limit(allocated_size)
-
-                move_weight_functions(m, device_to)
-
-            logging.info(f"Model {self.model.__class__.__name__} prepared for dynamic VRAM loading. {allocated_size // (1024 ** 2)}MB Staged. {num_patches} patches attached.")
-
-            self.model.device = device_to
-            self.model.current_weight_patches_uuid = self.patches_uuid
-
-            for callback in self.get_all_callbacks(CallbacksMP.ON_LOAD):
-                #These are all super dangerous. Who knows what the custom nodes actually do here...
-                callback(self, device_to, lowvram_model_memory, force_patch_weights, full_load)
-
-            self.apply_hooks(self.forced_hooks, force_apply=True)
-
-    def partially_unload(self, device_to, memory_to_free=0, force_patch_weights=False):
-        assert not force_patch_weights #See above
-        assert self.load_device != torch.device("cpu")
-
-        vbar = self._vbar_get()
-        return 0 if vbar is None else vbar.free_memory(memory_to_free)
-
-    def partially_unload_ram(self, ram_to_unload):
-        loading = self._load_list(prio_comfy_cast_weights=True, default_device=self.offload_device)
-        for x in loading:
-            _, _, _, _, m, _ = x
-            ram_to_unload -= comfy.pinned_memory.unpin_memory(m)
-            if ram_to_unload <= 0:
-                return
-
-    def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False):
-        #This isn't used by the core at all and can only be to load a model out of
-        #the control of proper model_managment. If you are a custom node author reading
-        #this, the correct pattern is to call load_models_gpu() to get a proper
-        #managed load of your model.
-        assert not load_weights
-        return super().patch_model(load_weights=load_weights, force_patch_weights=force_patch_weights)
-
-    def unpatch_model(self, device_to=None, unpatch_weights=True):
-        super().unpatch_model(device_to=None, unpatch_weights=False)
-
-        if unpatch_weights:
-            self.partially_unload_ram(1e32)
-            self.partially_unload(None, 1e32)
-            for m in self.model.modules():
-                move_weight_functions(m, device_to)
-
-            keys = list(self.backup.keys())
-            for k in keys:
-                bk = self.backup[k]
-                comfy.utils.set_attr_param(self.model, k, bk.weight)
-
-    def partially_load(self, device_to, extra_memory=0, force_patch_weights=False):
-        assert not force_patch_weights #See above
-        with self.use_ejected(skip_and_inject_on_exit_only=True):
-            dirty = self.model.current_weight_patches_uuid is not None and (self.model.current_weight_patches_uuid != self.patches_uuid)
-
-            self.unpatch_model(self.offload_device, unpatch_weights=False)
-            self.patch_model(load_weights=False)
-
-            try:
-                self.load(device_to, dirty=dirty)
-            except Exception as e:
-                self.detach()
-                raise e
-            #ModelPatcher::partially_load returns a number on what got loaded but
-            #nothing in core uses this and we have no data in the Dynamic world. Hit
-            #the custom node devs with a None rather than a 0 that would mislead any
-            #logic they might have.
-            return None
-
-    def patch_cached_hook_weights(self, cached_weights: dict, key: str, memory_counter: MemoryCounter):
-        assert False #Should be unreachable - we dont ever cache in the new implementation
-
-    def patch_hook_weight_to_device(self, hooks: comfy.hooks.HookGroup, combined_patches: dict, key: str, original_weights: dict, memory_counter: MemoryCounter):
-        if key not in combined_patches:
-            return
-
-        raise RuntimeError("Hooks not implemented in ModelPatcherDynamic. Please remove --fast arguments form ComfyUI startup")
-
-    def unpatch_hooks(self, whitelist_keys_set: set[str]=None) -> None:
-        pass
-
-CoreModelPatcher = ModelPatcher
@@ -19,15 +19,10 @@
 import torch
 import logging
 import comfy.model_management
-from comfy.cli_args import args, PerformanceFeature, enables_dynamic_vram
+from comfy.cli_args import args, PerformanceFeature
 import comfy.float
+import comfy.rmsnorm
 import json
-import comfy.memory_management
-import comfy.pinned_memory
-import comfy.utils
-
-import comfy_aimdo.model_vbar
-import comfy_aimdo.torch

 def run_every_op():
    if torch.compiler.is_compiling():
@@ -53,8 +48,6 @@ try:
            SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION)

            def scaled_dot_product_attention(q, k, v, *args, **kwargs):
-                if q.nelement() < 1024 * 128:  # arbitrary number, for small inputs cudnn attention seems slower
-                    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
                with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
                    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
        else:
@@ -79,122 +72,7 @@ def cast_to_input(weight, input, non_blocking=False, copy=True):
    return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)


-def cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compute_dtype, want_requant):
-    offload_stream = None
-    xfer_dest = None
-
-    signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
-    resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
-    if signature is not None:
-        if resident:
-            weight = s._v_weight
-            bias = s._v_bias
-        else:
-            xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device)
-
-    if not resident:
-        cast_geometry = comfy.memory_management.tensors_to_geometries([ s.weight, s.bias ])
-        cast_dest = None
-
-        xfer_source = [ s.weight, s.bias ]
-
-        pin = comfy.pinned_memory.get_pin(s)
-        if pin is not None:
-            xfer_source = [ pin ]
-
-        for data, geometry in zip([ s.weight, s.bias ], cast_geometry):
-            if data is None:
-                continue
-            if data.dtype != geometry.dtype:
-                cast_dest = xfer_dest
-                if cast_dest is None:
-                    cast_dest = torch.empty((comfy.memory_management.vram_aligned_size(cast_geometry),), dtype=torch.uint8, device=device)
-                xfer_dest = None
-                break
-
-        dest_size = comfy.memory_management.vram_aligned_size(xfer_source)
-        offload_stream = comfy.model_management.get_offload_stream(device)
-        if xfer_dest is None and offload_stream is not None:
-                xfer_dest = comfy.model_management.get_cast_buffer(offload_stream, device, dest_size, s)
-                if xfer_dest is None:
-                    offload_stream = comfy.model_management.get_offload_stream(device)
-                    xfer_dest = comfy.model_management.get_cast_buffer(offload_stream, device, dest_size, s)
-        if xfer_dest is None:
-            xfer_dest = torch.empty((dest_size,), dtype=torch.uint8, device=device)
-            offload_stream = None
-
-        if signature is None and pin is None:
-            comfy.pinned_memory.pin_memory(s)
-            pin = comfy.pinned_memory.get_pin(s)
-        else:
-            pin = None
-
-        if pin is not None:
-            comfy.model_management.cast_to_gathered(xfer_source, pin)
-            xfer_source = [ pin ]
-        #send it over
-        comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=offload_stream)
-        comfy.model_management.sync_stream(device, offload_stream)
-
-        if cast_dest is not None:
-            for pre_cast, post_cast in zip(comfy.memory_management.interpret_gathered_like([s.weight, s.bias ], xfer_dest),
-                                           comfy.memory_management.interpret_gathered_like(cast_geometry, cast_dest)):
-                if post_cast is not None:
-                    post_cast.copy_(pre_cast)
-            xfer_dest = cast_dest
-
-        params = comfy.memory_management.interpret_gathered_like(cast_geometry, xfer_dest)
-        weight = params[0]
-        bias = params[1]
-        if signature is not None:
-            s._v_weight = weight
-            s._v_bias = bias
-        s._v_signature=signature
-
-    def post_cast(s, param_key, x, dtype, resident, update_weight):
-        lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
-        fns = getattr(s, param_key + "_function", [])
-
-        orig = x
-
-        def to_dequant(tensor, dtype):
-            tensor = tensor.to(dtype=dtype)
-            if isinstance(tensor, QuantizedTensor):
-                tensor = tensor.dequantize()
-            return tensor
-
-        if orig.dtype != dtype or len(fns) > 0:
-            x = to_dequant(x, dtype)
-        if not resident and lowvram_fn is not None:
-            x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
-            #FIXME: this is not accurate, we need to be sensitive to the compute dtype
-            x = lowvram_fn(x)
-            if (isinstance(orig, QuantizedTensor) and
-                (want_requant and len(fns) == 0 or update_weight)):
-                seed = comfy.utils.string_to_seed(s.seed_key)
-                y = QuantizedTensor.from_float(x, s.layout_type, scale="recalculate", stochastic_rounding=seed)
-                if want_requant and len(fns) == 0:
-                    #The layer actually wants our freshly saved QT
-                    x = y
-            elif update_weight:
-                y = comfy.float.stochastic_rounding(x, orig.dtype, seed = comfy.utils.string_to_seed(s.seed_key))
-            if update_weight:
-                orig.copy_(y)
-        for f in fns:
-            x = f(x)
-        return x
-
-    update_weight = signature is not None
-
-    weight = post_cast(s, "weight", weight, dtype, resident, update_weight)
-    if s.bias is not None:
-        bias = post_cast(s, "bias", bias, bias_dtype, resident, update_weight)
-
-    #FIXME: weird offload return protocol
-    return weight, bias, (offload_stream, device if signature is not None else None, None)
-
-
-def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):
+def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False):
    # NOTE: offloadable=False is a a legacy and if you are a custom node author reading this please pass
    # offloadable=True and call uncast_bias_weight() after your last usage of the weight/bias. This
    # will add async-offload support to your cast and improve performance.
@@ -209,38 +87,22 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
        if device is None:
            device = input.device

-    non_blocking = comfy.model_management.device_supports_non_blocking(device)
-
-    if hasattr(s, "_v"):
-        return cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compute_dtype, want_requant)
-
    if offloadable and (device != s.weight.device or
                        (s.bias is not None and device != s.bias.device)):
        offload_stream = comfy.model_management.get_offload_stream(device)
    else:
        offload_stream = None

-    bias = None
-    weight = None
-
-    if offload_stream is not None and not args.cuda_malloc:
-        cast_buffer_size = comfy.memory_management.vram_aligned_size([ s.weight, s.bias ])
-        cast_buffer = comfy.model_management.get_cast_buffer(offload_stream, device, cast_buffer_size, s)
-        #The streams can be uneven in buffer capability and reject us. Retry to get the other stream
-        if cast_buffer is None:
-            offload_stream = comfy.model_management.get_offload_stream(device)
-            cast_buffer = comfy.model_management.get_cast_buffer(offload_stream, device, cast_buffer_size, s)
-        params = comfy.memory_management.interpret_gathered_like([ s.weight, s.bias ], cast_buffer)
-        weight = params[0]
-        bias = params[1]
+    non_blocking = comfy.model_management.device_supports_non_blocking(device)

    weight_has_function = len(s.weight_function) > 0
    bias_has_function = len(s.bias_function) > 0

-    weight = comfy.model_management.cast_to(s.weight, None, device, non_blocking=non_blocking, copy=weight_has_function, stream=offload_stream, r=weight)
+    weight = comfy.model_management.cast_to(s.weight, None, device, non_blocking=non_blocking, copy=weight_has_function, stream=offload_stream)

+    bias = None
    if s.bias is not None:
-        bias = comfy.model_management.cast_to(s.bias, None, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream, r=bias)
+        bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream)

    comfy.model_management.sync_stream(device, offload_stream)

@@ -248,7 +110,6 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
    weight_a = weight

    if s.bias is not None:
-        bias = bias.to(dtype=bias_dtype)
        for f in s.bias_function:
            bias = f(bias)

@@ -270,20 +131,14 @@ def uncast_bias_weight(s, weight, bias, offload_stream):
    if offload_stream is None:
        return
    os, weight_a, bias_a = offload_stream
-    device=None
-    #FIXME: This is not good RTTI
-    if not isinstance(weight_a, torch.Tensor):
-        comfy_aimdo.model_vbar.vbar_unpin(s._v)
-        device = weight_a
    if os is None:
        return
-    if device is None:
-        if weight_a is not None:
-            device = weight_a.device
-        else:
-            if bias_a is None:
-                return
-            device = bias_a.device
+    if weight_a is not None:
+        device = weight_a.device
+    else:
+        if bias_a is None:
+            return
+        device = bias_a.device
    os.wait_stream(comfy.model_management.current_stream(device))


@@ -294,57 +149,6 @@ class CastWeightBiasOp:

 class disable_weight_init:
    class Linear(torch.nn.Linear, CastWeightBiasOp):
-
-        def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
-            if not comfy.model_management.WINDOWS or not enables_dynamic_vram():
-                super().__init__(in_features, out_features, bias, device, dtype)
-                return
-
-            # Issue is with `torch.empty` still reserving the full memory for the layer.
-            # Windows doesn't over-commit memory so without this, We are momentarily commit
-            # charged for the weight even though we might zero-copy it when we load the
-            # state dict. If the commit charge exceeds the ceiling we can destabilize the
-            # system.
-            torch.nn.Module.__init__(self)
-            self.in_features = in_features
-            self.out_features = out_features
-            self.weight = None
-            self.bias = None
-            self.comfy_need_lazy_init_bias=bias
-            self.weight_comfy_model_dtype = dtype
-            self.bias_comfy_model_dtype = dtype
-
-        def _load_from_state_dict(self, state_dict, prefix, local_metadata,
-                                strict, missing_keys, unexpected_keys, error_msgs):
-
-            if not comfy.model_management.WINDOWS or not enables_dynamic_vram():
-                return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
-                                                     missing_keys, unexpected_keys, error_msgs)
-            assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)
-            prefix_len = len(prefix)
-            for k,v in state_dict.items():
-                if k[prefix_len:] == "weight":
-                    if not assign_to_params_buffers:
-                        v = v.clone()
-                    self.weight = torch.nn.Parameter(v, requires_grad=False)
-                elif k[prefix_len:] == "bias" and v is not None:
-                    if not assign_to_params_buffers:
-                        v = v.clone()
-                    self.bias = torch.nn.Parameter(v, requires_grad=False)
-                else:
-                    unexpected_keys.append(k)
-
-            #Reconcile default construction of the weight if its missing.
-            if self.weight is None:
-                v = torch.zeros(self.in_features, self.out_features)
-                self.weight = torch.nn.Parameter(v, requires_grad=False)
-                missing_keys.append(prefix+"weight")
-            if self.bias is None and self.comfy_need_lazy_init_bias:
-                v = torch.zeros(self.out_features,)
-                self.bias = torch.nn.Parameter(v, requires_grad=False)
-                missing_keys.append(prefix+"bias")
-
-
        def reset_parameters(self):
            return None

@@ -462,7 +266,7 @@ class disable_weight_init:
            else:
                return super().forward(*args, **kwargs)

-    class RMSNorm(torch.nn.RMSNorm, CastWeightBiasOp):
+    class RMSNorm(comfy.rmsnorm.RMSNorm, CastWeightBiasOp):
        def reset_parameters(self):
            self.bias = None
            return None
@@ -474,7 +278,8 @@ class disable_weight_init:
                weight = None
                bias = None
                offload_stream = None
-            x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
+            x = comfy.rmsnorm.rms_norm(input, weight, self.eps)  # TODO: switch to commented out line when old torch is deprecated
+            # x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
            uncast_bias_weight(self, weight, bias, offload_stream)
            return x

@@ -850,8 +655,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
            def _forward(self, input, weight, bias):
                return torch.nn.functional.linear(input, weight, bias)

-            def forward_comfy_cast_weights(self, input, compute_dtype=None, want_requant=False):
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True, compute_dtype=compute_dtype, want_requant=want_requant)
+            def forward_comfy_cast_weights(self, input):
+                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
                x = self._forward(input, weight, bias)
                uncast_bias_weight(self, weight, bias, offload_stream)
                return x
@@ -861,8 +666,6 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec

                input_shape = input.shape
                reshaped_3d = False
-                #If cast needs to apply lora, it should be done in the compute dtype
-                compute_dtype = input.dtype

                if (getattr(self, 'layout_type', None) is not None and
                    not isinstance(input, QuantizedTensor) and not self._full_precision_mm and
@@ -881,7 +684,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                            scale = comfy.model_management.cast_to_device(scale, input.device, None)
                        input = QuantizedTensor.from_float(input_reshaped, self.layout_type, scale=scale)

-                output = self.forward_comfy_cast_weights(input, compute_dtype, want_requant=isinstance(input, QuantizedTensor))
+                output = self.forward_comfy_cast_weights(input)

                # Reshape output back to 3D if input was 3D
                if reshaped_3d:
@@ -1,29 +0,0 @@
-import torch
-import comfy.model_management
-import comfy.memory_management
-
-from comfy.cli_args import args
-
-def get_pin(module):
-    return getattr(module, "_pin", None)
-
-def pin_memory(module):
-    if module.pin_failed or args.disable_pinned_memory or get_pin(module) is not None:
-        return
-    #FIXME: This is a RAM cache trigger event
-    size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
-    pin = torch.empty((size,), dtype=torch.uint8)
-    if comfy.model_management.pin_memory(pin):
-        module._pin = pin
-    else:
-        module.pin_failed = True
-        return False
-    return True
-
-def unpin_memory(module):
-    if get_pin(module) is None:
-        return 0
-    size = module._pin.numel() * module._pin.element_size()
-    comfy.model_management.unpin_memory(module._pin)
-    del module._pin
-    return size
@@ -1,10 +1,57 @@
 import torch
 import comfy.model_management
+import numbers
+import logging
+
+RMSNorm = None
+
+try:
+    rms_norm_torch = torch.nn.functional.rms_norm
+    RMSNorm = torch.nn.RMSNorm
+except:
+    rms_norm_torch = None
+    logging.warning("Please update pytorch to use native RMSNorm")

-RMSNorm = torch.nn.RMSNorm

 def rms_norm(x, weight=None, eps=1e-6):
-    if weight is None:
-        return torch.nn.functional.rms_norm(x, (x.shape[-1],), eps=eps)
+    if rms_norm_torch is not None and not (torch.jit.is_tracing() or torch.jit.is_scripting()):
+        if weight is None:
+            return rms_norm_torch(x, (x.shape[-1],), eps=eps)
+        else:
+            return rms_norm_torch(x, weight.shape, weight=comfy.model_management.cast_to(weight, dtype=x.dtype, device=x.device), eps=eps)
    else:
-        return torch.nn.functional.rms_norm(x, weight.shape, weight=comfy.model_management.cast_to(weight, dtype=x.dtype, device=x.device), eps=eps)
+        r = x * torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + eps)
+        if weight is None:
+            return r
+        else:
+            return r * comfy.model_management.cast_to(weight, dtype=x.dtype, device=x.device)
+
+
+if RMSNorm is None:
+    class RMSNorm(torch.nn.Module):
+        def __init__(
+            self,
+            normalized_shape,
+            eps=1e-6,
+            elementwise_affine=True,
+            device=None,
+            dtype=None,
+        ):
+            factory_kwargs = {"device": device, "dtype": dtype}
+            super().__init__()
+            if isinstance(normalized_shape, numbers.Integral):
+                # mypy error: incompatible types in assignment
+                normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+            self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+            self.eps = eps
+            self.elementwise_affine = elementwise_affine
+            if self.elementwise_affine:
+                self.weight = torch.nn.Parameter(
+                    torch.empty(self.normalized_shape, **factory_kwargs)
+                )
+            else:
+                self.register_parameter("weight", None)
+            self.bias = None
+
+        def forward(self, x):
+            return rms_norm(x, self.weight, self.eps)
@@ -122,26 +122,20 @@ def estimate_memory(model, noise_shape, conds):
    minimum_memory_required = model.model.memory_required([noise_shape[0]] + list(noise_shape[1:]), cond_shapes=cond_shapes_min)
    return memory_required, minimum_memory_required

-def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None, force_full_load=False, force_offload=False):
+def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None, force_full_load=False):
    executor = comfy.patcher_extension.WrapperExecutor.new_executor(
        _prepare_sampling,
        comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING, model_options, is_model_options=True)
    )
-    return executor.execute(model, noise_shape, conds, model_options=model_options, force_full_load=force_full_load, force_offload=force_offload)
+    return executor.execute(model, noise_shape, conds, model_options=model_options, force_full_load=force_full_load)

-def _prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None, force_full_load=False, force_offload=False):
+def _prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None, force_full_load=False):
    real_model: BaseModel = None
    models, inference_memory = get_additional_models(conds, model.model_dtype())
    models += get_additional_models_from_model_options(model_options)
    models += model.get_nested_additional_models()  # TODO: does this require inference_memory update?
-    if force_offload: # In training + offload enabled, we want to force prepare sampling to trigger partial load
-        memory_required = 1e20
-        minimum_memory_required = None
-    else:
-        memory_required, minimum_memory_required = estimate_memory(model, noise_shape, conds)
-        memory_required += inference_memory
-        minimum_memory_required += inference_memory
-    comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required, minimum_memory_required=minimum_memory_required, force_full_load=force_full_load)
+    memory_required, minimum_memory_required = estimate_memory(model, noise_shape, conds)
+    comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required + inference_memory, minimum_memory_required=minimum_memory_required + inference_memory, force_full_load=force_full_load)
    real_model = model.model

    return real_model, conds, models
@@ -9,6 +9,7 @@ if TYPE_CHECKING:
 import torch
 from functools import partial
 import collections
+from comfy import model_management
 import math
 import logging
 import comfy.sampler_helpers
@@ -259,7 +260,7 @@ def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tens
            to_batch_temp.reverse()
            to_batch = to_batch_temp[:1]

-            free_memory = model.current_patcher.get_free_memory(x_in.device)
+            free_memory = model_management.get_free_memory(x_in.device)
            for i in range(1, len(to_batch_temp) + 1):
                batch_amount = to_batch_temp[:len(to_batch_temp)//i]
                input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:]
@@ -59,7 +59,6 @@ import comfy.text_encoders.kandinsky5
 import comfy.text_encoders.jina_clip_2
 import comfy.text_encoders.newbie
 import comfy.text_encoders.anima
-import comfy.text_encoders.ace15

 import comfy.model_patcher
 import comfy.lora
@@ -229,10 +228,8 @@ class CLIP:
                    self.cond_stage_model.to(offload_device)
                    logging.warning("Had to shift TE back.")

-        model_management.archive_model_dtypes(self.cond_stage_model)
-
        self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-        self.patcher = comfy.model_patcher.CoreModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
+        self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
        #Match torch.float32 hardcode upcast in TE implemention
        self.patcher.set_model_compute_dtype(torch.float32)
        self.patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram
@@ -392,18 +389,8 @@ class CLIP:

    def load_sd(self, sd, full_model=False):
        if full_model:
-            return self.cond_stage_model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
+            return self.cond_stage_model.load_state_dict(sd, strict=False)
        else:
-            can_assign = self.patcher.is_dynamic()
-            self.cond_stage_model.can_assign_sd = can_assign
-
-            # The CLIP models are a pretty complex web of wrappers and its
-            # a bit of an API change to plumb this all the way through.
-            # So spray paint the model with this flag that the loading
-            # nn.Module can then inspect for itself.
-            for m in self.cond_stage_model.modules():
-                m.can_assign_sd = can_assign
-
            return self.cond_stage_model.load_sd(sd)

    def get_sd(self):
@@ -423,17 +410,6 @@ class CLIP:
    def get_key_patches(self):
        return self.patcher.get_key_patches()

-    def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, min_p=0.0, repetition_penalty=1.0, seed=None):
-        self.cond_stage_model.reset_clip_options()
-
-        self.load_model()
-        self.cond_stage_model.set_clip_options({"layer": None})
-        self.cond_stage_model.set_clip_options({"execution_device": self.patcher.load_device})
-        return self.cond_stage_model.generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed)
-
-    def decode(self, token_ids, skip_special_tokens=True):
-        return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
-
 class VAE:
    def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
        if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
@@ -464,8 +440,6 @@ class VAE:
        self.extra_1d_channel = None
        self.crop_input = True

-        self.audio_sample_rate = 44100
-
        if config is None:
            if "decoder.mid.block_1.mix_factor" in sd:
                encoder_config = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
@@ -563,27 +537,14 @@ class VAE:
                                                                    encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
                                                                    decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})
            elif "decoder.layers.1.layers.0.beta" in sd:
-                config = {}
-                param_key = None
-                self.upscale_ratio = 2048
-                self.downscale_ratio = 2048
-                if "decoder.layers.2.layers.1.weight_v" in sd:
-                    param_key = "decoder.layers.2.layers.1.weight_v"
-                if "decoder.layers.2.layers.1.parametrizations.weight.original1" in sd:
-                    param_key = "decoder.layers.2.layers.1.parametrizations.weight.original1"
-                if param_key is not None:
-                    if sd[param_key].shape[-1] == 12:
-                        config["strides"] = [2, 4, 4, 6, 10]
-                        self.audio_sample_rate = 48000
-                        self.upscale_ratio = 1920
-                        self.downscale_ratio = 1920
-
-                self.first_stage_model = AudioOobleckVAE(**config)
+                self.first_stage_model = AudioOobleckVAE()
                self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)
                self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * 2048) * model_management.dtype_size(dtype)
                self.latent_channels = 64
                self.output_channels = 2
                self.pad_channel_value = "replicate"
+                self.upscale_ratio = 2048
+                self.downscale_ratio =  2048
                self.latent_dim = 1
                self.process_output = lambda audio: audio
                self.process_input = lambda audio: audio
@@ -804,6 +765,13 @@ class VAE:
            self.first_stage_model = AutoencoderKL(**(config['params']))
        self.first_stage_model = self.first_stage_model.eval()

+        m, u = self.first_stage_model.load_state_dict(sd, strict=False)
+        if len(m) > 0:
+            logging.warning("Missing VAE keys {}".format(m))
+
+        if len(u) > 0:
+            logging.debug("Leftover VAE keys {}".format(u))
+
        if device is None:
            device = model_management.vae_device()
        self.device = device
@@ -812,21 +780,9 @@ class VAE:
            dtype = model_management.vae_dtype(self.device, self.working_dtypes)
        self.vae_dtype = dtype
        self.first_stage_model.to(self.vae_dtype)
-        model_management.archive_model_dtypes(self.first_stage_model)
        self.output_device = model_management.intermediate_device()

-        mp = comfy.model_patcher.CoreModelPatcher
-        if self.disable_offload:
-            mp = comfy.model_patcher.ModelPatcher
-        self.patcher = mp(self.first_stage_model, load_device=self.device, offload_device=offload_device)
-
-        m, u = self.first_stage_model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
-        if len(m) > 0:
-            logging.warning("Missing VAE keys {}".format(m))
-
-        if len(u) > 0:
-            logging.debug("Leftover VAE keys {}".format(u))
-
+        self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device)
        logging.info("VAE load device: {}, offload device: {}, dtype: {}".format(self.device, offload_device, self.vae_dtype))
        self.model_size()

@@ -882,7 +838,7 @@ class VAE:
            / 3.0)
        return output

-    def decode_tiled_1d(self, samples, tile_x=256, overlap=32):
+    def decode_tiled_1d(self, samples, tile_x=128, overlap=32):
        if samples.ndim == 3:
            decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
        else:
@@ -941,7 +897,7 @@ class VAE:
        try:
            memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
            model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
-            free_memory = self.patcher.get_free_memory(self.device)
+            free_memory = model_management.get_free_memory(self.device)
            batch_number = int(free_memory / memory_used)
            batch_number = max(1, batch_number)

@@ -986,7 +942,7 @@ class VAE:
        if overlap is not None:
            args["overlap"] = overlap

-        if dims == 1 or self.extra_1d_channel is not None:
+        if dims == 1:
            args.pop("tile_y")
            output = self.decode_tiled_1d(samples, **args)
        elif dims == 2:
@@ -1015,7 +971,7 @@ class VAE:
        try:
            memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
            model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
-            free_memory = self.patcher.get_free_memory(self.device)
+            free_memory = model_management.get_free_memory(self.device)
            batch_number = int(free_memory / max(1, memory_used))
            batch_number = max(1, batch_number)
            samples = None
@@ -1193,7 +1149,6 @@ class TEModel(Enum):
    JINA_CLIP_2 = 19
    QWEN3_8B = 20
    QWEN3_06B = 21
-    GEMMA_3_4B_VISION = 22


 def detect_te_model(sd):
@@ -1222,10 +1177,7 @@ def detect_te_model(sd):
        if 'model.layers.47.self_attn.q_norm.weight' in sd:
            return TEModel.GEMMA_3_12B
        if 'model.layers.0.self_attn.q_norm.weight' in sd:
-            if 'vision_model.embeddings.patch_embedding.weight' in sd:
-                return TEModel.GEMMA_3_4B_VISION
-            else:
-                return TEModel.GEMMA_3_4B
+            return TEModel.GEMMA_3_4B
        return TEModel.GEMMA_2_2B
    if 'model.layers.0.self_attn.k_proj.bias' in sd:
        weight = sd['model.layers.0.self_attn.k_proj.bias']
@@ -1285,8 +1237,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
        else:
            if "text_projection" in clip_data[i]:
                clip_data[i]["text_projection.weight"] = clip_data[i]["text_projection"].transpose(0, 1) #old models saved with the CLIPSave node
-        if "lm_head.weight" in clip_data[i]:
-            clip_data[i]["model.lm_head.weight"] = clip_data[i].pop("lm_head.weight") # prefix missing in some models

    tokenizer_data = {}
    clip_target = EmptyClass()
@@ -1352,14 +1302,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
            clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data), model_type="gemma3_4b")
            clip_target.tokenizer = comfy.text_encoders.lumina2.NTokenizer
            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
-        elif te_model == TEModel.GEMMA_3_4B_VISION:
-            clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data), model_type="gemma3_4b_vision")
-            clip_target.tokenizer = comfy.text_encoders.lumina2.NTokenizer
-            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
-        elif te_model == TEModel.GEMMA_3_12B:
-            clip_target.clip = comfy.text_encoders.lt.gemma3_te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.lt.Gemma3_12BTokenizer
-            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
        elif te_model == TEModel.LLAMA3_8:
            clip_target.clip = comfy.text_encoders.hidream.hidream_clip(**llama_detect(clip_data),
                                                                        clip_l=False, clip_g=False, t5=False, llama=True, dtype_t5=None)
@@ -1467,14 +1409,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
                clip_data_jina = clip_data[0]
            tokenizer_data["gemma_spiece_model"] = clip_data_gemma.get("spiece_model", None)
            tokenizer_data["jina_spiece_model"] = clip_data_jina.get("spiece_model", None)
-        elif clip_type == CLIPType.ACE:
-            te_models = [detect_te_model(clip_data[0]), detect_te_model(clip_data[1])]
-            if TEModel.QWEN3_4B in te_models:
-                model_type = "qwen3_4b"
-            else:
-                model_type = "qwen3_2b"
-            clip_target.clip = comfy.text_encoders.ace15.te(lm_model=model_type, **llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.ace15.ACE15Tokenizer
        else:
            clip_target.clip = sdxl_clip.SDXLClipModel
            clip_target.tokenizer = sdxl_clip.SDXLTokenizer
@@ -1498,7 +1432,7 @@ def load_gligen(ckpt_path):
    model = gligen.load_gligen(data)
    if model_management.should_use_fp16():
        model = model.half()
-    return comfy.model_patcher.CoreModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
+    return comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())

 def model_detection_error_hint(path, state_dict):
    filename = os.path.basename(path)
@@ -1586,8 +1520,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
    if output_model:
        inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)
        model = model_config.get_model(sd, diffusion_model_prefix, device=inital_load_device)
-        model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
-        model.load_model_weights(sd, diffusion_model_prefix, assign=model_patcher.is_dynamic())
+        model.load_model_weights(sd, diffusion_model_prefix)

    if output_vae:
        vae_sd = comfy.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True)
@@ -1630,6 +1563,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
        logging.debug("left over keys: {}".format(left_over))

    if output_model:
+        model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
        if inital_load_device != torch.device("cpu"):
            logging.info("loaded diffusion model directly to GPU")
            model_management.load_models_gpu([model_patcher], force_full_load=True)
@@ -1721,14 +1655,13 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None):
        model_config.optimizations["fp8"] = True

    model = model_config.get_model(new_sd, "")
-    model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=load_device, offload_device=offload_device)
-    if not model_management.is_device_cpu(offload_device):
-        model.to(offload_device)
-    model.load_model_weights(new_sd, "", assign=model_patcher.is_dynamic())
+    model = model.to(offload_device)
+    model.load_model_weights(new_sd, "")
    left_over = sd.keys()
    if len(left_over) > 0:
        logging.info("left over keys in diffusion model: {}".format(left_over))
-    return model_patcher
+    return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
+

 def load_diffusion_model(unet_path, model_options={}):
    sd, metadata = comfy.utils.load_torch_file(unet_path, return_metadata=True)
@@ -1759,9 +1692,9 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m
    if metadata is None:
        metadata = {}

-    model_management.load_models_gpu(load_models)
+    model_management.load_models_gpu(load_models, force_patch_weights=True)
    clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None
-    sd = model.state_dict_for_saving(clip_sd, vae_sd, clip_vision_sd)
+    sd = model.model.state_dict_for_saving(clip_sd, vae_sd, clip_vision_sd)
    for k in extra_keys:
        sd[k] = extra_keys[k]

@@ -155,8 +155,6 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
        self.execution_device = options.get("execution_device", self.execution_device)
        if isinstance(self.layer, list) or self.layer == "all":
            pass
-        elif isinstance(layer_idx, list):
-            self.layer = layer_idx
        elif layer_idx is None or abs(layer_idx) > self.num_layers:
            self.layer = "last"
        else:
@@ -171,9 +169,8 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):

    def process_tokens(self, tokens, device):
        end_token = self.special_tokens.get("end", None)
-        pad_token = self.special_tokens.get("pad", -1)
        if end_token is None:
-            cmp_token = pad_token
+            cmp_token = self.special_tokens.get("pad", -1)
        else:
            cmp_token = end_token

@@ -187,21 +184,15 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
            other_embeds = []
            eos = False
            index = 0
-            left_pad = False
            for y in x:
                if isinstance(y, numbers.Integral):
-                    token = int(y)
-                    if index == 0 and token == pad_token:
-                        left_pad = True
-
-                    if eos or (left_pad and token == pad_token):
+                    if eos:
                        attention_mask.append(0)
                    else:
                        attention_mask.append(1)
-                        left_pad = False
-
+                    token = int(y)
                    tokens_temp += [token]
-                    if not eos and token == cmp_token and not left_pad:
+                    if not eos and token == cmp_token:
                        if end_token is None:
                            attention_mask[-1] = 0
                        eos = True
@@ -306,16 +297,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
        return self(tokens)

    def load_sd(self, sd):
-        return self.transformer.load_state_dict(sd, strict=False, assign=getattr(self, "can_assign_sd", False))
-
-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed):
-        if isinstance(tokens, dict):
-            tokens_only = next(iter(tokens.values())) # todo: get this better?
-        else:
-            tokens_only = tokens
-        tokens_only = [[t[0] for t in b] for b in tokens_only]
-        embeds = self.process_tokens(tokens_only, device=self.execution_device)[0]
-        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed)
+        return self.transformer.load_state_dict(sd, strict=False)

 def parse_parentheses(string):
    result = []
@@ -573,8 +555,6 @@ class SDTokenizer:
        min_length = tokenizer_options.get("{}_min_length".format(self.embedding_key), self.min_length)
        min_padding = tokenizer_options.get("{}_min_padding".format(self.embedding_key), self.min_padding)

-        min_length = kwargs.get("min_length", min_length)
-
        text = escape_important(text)
        if kwargs.get("disable_weights", self.disable_weights):
            parsed_weights = [(text, 1.0)]
@@ -674,9 +654,6 @@ class SDTokenizer:
    def state_dict(self):
        return {}

-    def decode(self, token_ids, skip_special_tokens=True):
-        return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
-
 class SD1Tokenizer:
    def __init__(self, embedding_directory=None, tokenizer_data={}, clip_name="l", tokenizer=SDTokenizer, name=None):
        if name is not None:
@@ -700,9 +677,6 @@ class SD1Tokenizer:
    def state_dict(self):
        return getattr(self, self.clip).state_dict()

-    def decode(self, token_ids, skip_special_tokens=True):
-        return getattr(self, self.clip).decode(token_ids, skip_special_tokens=skip_special_tokens)
-
 class SD1CheckpointClipModel(SDClipModel):
    def __init__(self, device="cpu", dtype=None, model_options={}):
        super().__init__(device=device, return_projected_pooled=False, dtype=dtype, model_options=model_options)
@@ -739,6 +713,3 @@ class SD1ClipModel(torch.nn.Module):

    def load_sd(self, sd):
        return getattr(self, self.clip).load_sd(sd)
-
-    def generate(self, tokens, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.95, min_p=0.0, repetition_penalty=1.0, seed=None):
-        return getattr(self, self.clip).generate(tokens, do_sample=do_sample, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, repetition_penalty=repetition_penalty, seed=seed)
@@ -24,7 +24,6 @@ import comfy.text_encoders.hunyuan_image
 import comfy.text_encoders.kandinsky5
 import comfy.text_encoders.z_image
 import comfy.text_encoders.anima
-import comfy.text_encoders.ace15

 from . import supported_models_base
 from . import latent_formats
@@ -710,15 +709,6 @@ class Flux(supported_models_base.BASE):

    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]

-    def process_unet_state_dict(self, state_dict):
-        out_sd = {}
-        for k in list(state_dict.keys()):
-            key_out = k
-            if key_out.endswith("_norm.scale"):
-                key_out = "{}.weight".format(key_out[:-len(".scale")])
-            out_sd[key_out] = state_dict[k]
-        return out_sd
-
    vae_key_prefix = ["vae."]
    text_encoder_key_prefix = ["text_encoders."]

@@ -907,13 +897,11 @@ class HunyuanVideo(supported_models_base.BASE):
            key_out = key_out.replace("txt_in.c_embedder.linear_1.", "txt_in.c_embedder.in_layer.").replace("txt_in.c_embedder.linear_2.", "txt_in.c_embedder.out_layer.")
            key_out = key_out.replace("_mod.linear.", "_mod.lin.").replace("_attn_qkv.", "_attn.qkv.")
            key_out = key_out.replace("mlp.fc1.", "mlp.0.").replace("mlp.fc2.", "mlp.2.")
-            key_out = key_out.replace("_attn_q_norm.weight", "_attn.norm.query_norm.weight").replace("_attn_k_norm.weight", "_attn.norm.key_norm.weight")
-            key_out = key_out.replace(".q_norm.weight", ".norm.query_norm.weight").replace(".k_norm.weight", ".norm.key_norm.weight")
+            key_out = key_out.replace("_attn_q_norm.weight", "_attn.norm.query_norm.scale").replace("_attn_k_norm.weight", "_attn.norm.key_norm.scale")
+            key_out = key_out.replace(".q_norm.weight", ".norm.query_norm.scale").replace(".k_norm.weight", ".norm.key_norm.scale")
            key_out = key_out.replace("_attn_proj.", "_attn.proj.")
            key_out = key_out.replace(".modulation.linear.", ".modulation.lin.")
            key_out = key_out.replace("_in.mlp.2.", "_in.out_layer.").replace("_in.mlp.0.", "_in.in_layer.")
-            if key_out.endswith(".scale"):
-                key_out = "{}.weight".format(key_out[:-len(".scale")])
            out_sd[key_out] = state_dict[k]
        return out_sd

@@ -1004,7 +992,7 @@ class CosmosT2IPredict2(supported_models_base.BASE):

    memory_usage_factor = 1.0

-    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]

    def __init__(self, unet_config):
        super().__init__(unet_config)
@@ -1034,7 +1022,11 @@ class Anima(supported_models_base.BASE):

    memory_usage_factor = 1.0

-    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    def __init__(self, unet_config):
+        super().__init__(unet_config)
+        self.memory_usage_factor = (unet_config.get("model_channels", 2048) / 2048) * 0.95

    def get_model(self, state_dict, prefix="", device=None):
        out = model_base.Anima(self, device=device)
@@ -1045,12 +1037,6 @@ class Anima(supported_models_base.BASE):
        detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3_06b.transformer.".format(pref))
        return supported_models_base.ClipTarget(comfy.text_encoders.anima.AnimaTokenizer, comfy.text_encoders.anima.te(**detect))

-    def set_inference_dtype(self, dtype, manual_cast_dtype, **kwargs):
-        self.memory_usage_factor = (self.unet_config.get("model_channels", 2048) / 2048) * 0.95
-        if dtype is torch.float16:
-            self.memory_usage_factor *= 1.4
-        return super().set_inference_dtype(dtype, manual_cast_dtype, **kwargs)
-
 class CosmosI2VPredict2(CosmosT2IPredict2):
    unet_config = {
        "image_model": "cosmos_predict2",
@@ -1275,15 +1261,6 @@ class Hunyuan3Dv2(supported_models_base.BASE):

    latent_format = latent_formats.Hunyuan3Dv2

-    def process_unet_state_dict(self, state_dict):
-        out_sd = {}
-        for k in list(state_dict.keys()):
-            key_out = k
-            if key_out.endswith(".scale"):
-                key_out = "{}.weight".format(key_out[:-len(".scale")])
-            out_sd[key_out] = state_dict[k]
-        return out_sd
-
    def process_unet_state_dict_for_saving(self, state_dict):
        replace_prefix = {"": "model."}
        return utils.state_dict_prefix_replace(state_dict, replace_prefix)
@@ -1361,14 +1338,6 @@ class Chroma(supported_models_base.BASE):

    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]

-    def process_unet_state_dict(self, state_dict):
-        out_sd = {}
-        for k in list(state_dict.keys()):
-            key_out = k
-            if key_out.endswith(".scale"):
-                key_out = "{}.weight".format(key_out[:-len(".scale")])
-            out_sd[key_out] = state_dict[k]
-        return out_sd

    def get_model(self, state_dict, prefix="", device=None):
        out = model_base.Chroma(self, device=device)
@@ -1627,46 +1596,6 @@ class Kandinsky5Image(Kandinsky5):
        return supported_models_base.ClipTarget(comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage, comfy.text_encoders.kandinsky5.te(**hunyuan_detect))


-class ACEStep15(supported_models_base.BASE):
-    unet_config = {
-        "audio_model": "ace1.5",
-    }
-
-    unet_extra_config = {
-    }
-
-    sampling_settings = {
-        "multiplier": 1.0,
-        "shift": 3.0,
-    }
-
-    latent_format = comfy.latent_formats.ACEAudio15
-
-    memory_usage_factor = 4.7
-
-    supported_inference_dtypes = [torch.bfloat16, torch.float32]
-
-    vae_key_prefix = ["vae."]
-    text_encoder_key_prefix = ["text_encoders."]
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.ACEStep15(self, device=device)
-        return out
-
-    def clip_target(self, state_dict={}):
-        pref = self.text_encoder_key_prefix[0]
-        detect_2b = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3_2b.transformer.".format(pref))
-        detect_4b = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3_4b.transformer.".format(pref))
-        if "dtype_llama" in detect_2b:
-            detect = detect_2b
-            detect["lm_model"] = "qwen3_2b"
-        elif "dtype_llama" in detect_4b:
-            detect = detect_4b
-            detect["lm_model"] = "qwen3_4b"
-
-        return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
-
-
-models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, LTXAV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, ACEStep15, Omnigen2, QwenImage, Flux2, Kandinsky5Image, Kandinsky5, Anima]
+models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, LTXAV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage, Flux2, Kandinsky5Image, Kandinsky5, Anima]

 models += [SVD_img2vid]
@@ -1,348 +0,0 @@
-from .anima import Qwen3Tokenizer
-import comfy.text_encoders.llama
-from comfy import sd1_clip
-import torch
-import math
-import yaml
-import comfy.utils
-
-
-def sample_manual_loop_no_classes(
-    model,
-    ids=None,
-    execution_dtype=None,
-    cfg_scale: float = 2.0,
-    temperature: float = 0.85,
-    top_p: float = 0.9,
-    top_k: int = None,
-    min_p: float = 0.000,
-    seed: int = 1,
-    min_tokens: int = 1,
-    max_new_tokens: int = 2048,
-    audio_start_id: int = 151669,  # The cutoff ID for audio codes
-    audio_end_id: int = 215669,
-    eos_token_id: int = 151645,
-):
-    if ids is None:
-        return []
-    device = model.execution_device
-
-    if execution_dtype is None:
-        if comfy.model_management.should_use_bf16(device):
-            execution_dtype = torch.bfloat16
-        else:
-            execution_dtype = torch.float32
-
-    embeds, attention_mask, num_tokens, embeds_info = model.process_tokens(ids, device)
-    embeds_batch = embeds.shape[0]
-
-    output_audio_codes = []
-    past_key_values = []
-    generator = torch.Generator(device=device)
-    generator.manual_seed(seed)
-    model_config = model.transformer.model.config
-    past_kv_shape = [embeds_batch, model_config.num_key_value_heads, embeds.shape[1] + min_tokens, model_config.head_dim]
-
-    for x in range(model_config.num_hidden_layers):
-        past_key_values.append((torch.empty(past_kv_shape, device=device, dtype=execution_dtype), torch.empty(past_kv_shape, device=device, dtype=execution_dtype), 0))
-
-    progress_bar = comfy.utils.ProgressBar(max_new_tokens)
-
-    for step in comfy.utils.model_trange(max_new_tokens, desc="LM sampling"):
-        outputs = model.transformer(None, attention_mask, embeds=embeds.to(execution_dtype), num_tokens=num_tokens, intermediate_output=None, dtype=execution_dtype, embeds_info=embeds_info, past_key_values=past_key_values)
-        next_token_logits = model.transformer.logits(outputs[0])[:, -1]
-        past_key_values = outputs[2]
-
-        if cfg_scale != 1.0:
-            cond_logits = next_token_logits[0:1]
-            uncond_logits = next_token_logits[1:2]
-            cfg_logits = uncond_logits + cfg_scale * (cond_logits - uncond_logits)
-        else:
-            cfg_logits = next_token_logits[0:1]
-
-        use_eos_score = eos_token_id is not None and eos_token_id < audio_start_id and min_tokens < step
-        if use_eos_score:
-            eos_score = cfg_logits[:, eos_token_id].clone()
-
-        remove_logit_value = torch.finfo(cfg_logits.dtype).min
-        # Only generate audio tokens
-        cfg_logits[:, :audio_start_id] = remove_logit_value
-        cfg_logits[:, audio_end_id:] = remove_logit_value
-
-        if use_eos_score:
-            cfg_logits[:, eos_token_id] = eos_score
-
-        if top_k is not None and top_k > 0:
-            top_k_vals, _ = torch.topk(cfg_logits, top_k)
-            min_val = top_k_vals[..., -1, None]
-            cfg_logits[cfg_logits < min_val] = remove_logit_value
-
-        if min_p is not None and min_p > 0:
-            probs = torch.softmax(cfg_logits, dim=-1)
-            p_max = probs.max(dim=-1, keepdim=True).values
-            indices_to_remove = probs < (min_p * p_max)
-            cfg_logits[indices_to_remove] = remove_logit_value
-
-        if top_p is not None and top_p < 1.0:
-            sorted_logits, sorted_indices = torch.sort(cfg_logits, descending=True)
-            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-            sorted_indices_to_remove = cumulative_probs > top_p
-            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-            sorted_indices_to_remove[..., 0] = 0
-            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-            cfg_logits[indices_to_remove] = remove_logit_value
-
-        if temperature > 0:
-            cfg_logits = cfg_logits / temperature
-            next_token = torch.multinomial(torch.softmax(cfg_logits, dim=-1), num_samples=1, generator=generator).squeeze(1)
-        else:
-            next_token = torch.argmax(cfg_logits, dim=-1)
-
-        token = next_token.item()
-
-        if token == eos_token_id:
-            break
-
-        embed, _, _, _ = model.process_tokens([[token]], device)
-        embeds = embed.repeat(embeds_batch, 1, 1)
-        attention_mask = torch.cat([attention_mask, torch.ones((embeds_batch, 1), device=device, dtype=attention_mask.dtype)], dim=1)
-
-        output_audio_codes.append(token - audio_start_id)
-        progress_bar.update_absolute(step)
-
-    return output_audio_codes
-
-
-def generate_audio_codes(model, positive, negative, min_tokens=1, max_tokens=1024, seed=0, cfg_scale=2.0, temperature=0.85, top_p=0.9, top_k=0, min_p=0.000):
-    positive = [[token for token, _ in inner_list] for inner_list in positive]
-    positive = positive[0]
-
-    if cfg_scale != 1.0:
-        negative = [[token for token, _ in inner_list] for inner_list in negative]
-        negative = negative[0]
-
-        neg_pad = 0
-        if len(negative) < len(positive):
-            neg_pad = (len(positive) - len(negative))
-            negative = [model.special_tokens["pad"]] * neg_pad + negative
-
-        pos_pad = 0
-        if len(negative) > len(positive):
-            pos_pad = (len(negative) - len(positive))
-            positive = [model.special_tokens["pad"]] * pos_pad + positive
-
-        ids = [positive, negative]
-    else:
-        ids = [positive]
-
-    return sample_manual_loop_no_classes(model, ids, cfg_scale=cfg_scale, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, seed=seed, min_tokens=min_tokens, max_new_tokens=max_tokens)
-
-
-class ACE15Tokenizer(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="qwen3_06b", tokenizer=Qwen3Tokenizer)
-
-    def _metas_to_cot(self, *, return_yaml: bool = False, **kwargs) -> str:
-        user_metas = {
-            k: kwargs.pop(k)
-            for k in ("bpm", "duration", "keyscale", "timesignature")
-            if k in kwargs
-        }
-        timesignature = user_metas.get("timesignature")
-        if isinstance(timesignature, str) and timesignature.endswith("/4"):
-            user_metas["timesignature"] = timesignature[:-2]
-        user_metas = {
-            k: v if not isinstance(v, str) or not v.isdigit() else int(v)
-            for k, v in user_metas.items()
-            if v not in {"unspecified", None}
-        }
-        if len(user_metas):
-            meta_yaml = yaml.dump(user_metas, allow_unicode=True, sort_keys=True).strip()
-        else:
-            meta_yaml = ""
-        return f"<think>\n{meta_yaml}\n</think>" if not return_yaml else meta_yaml
-
-    def _metas_to_cap(self, **kwargs) -> str:
-        use_keys = ("bpm", "timesignature", "keyscale", "duration")
-        user_metas = { k: kwargs.pop(k, "N/A") for k in use_keys }
-        timesignature = user_metas.get("timesignature")
-        if isinstance(timesignature, str) and timesignature.endswith("/4"):
-            user_metas["timesignature"] = timesignature[:-2]
-        duration = user_metas["duration"]
-        if duration == "N/A":
-            user_metas["duration"] = "30 seconds"
-        elif isinstance(duration, (str, int, float)):
-            user_metas["duration"] = f"{math.ceil(float(duration))} seconds"
-        else:
-            raise TypeError("Unexpected type for duration key, must be str, int or float")
-        return "\n".join(f"- {k}: {user_metas[k]}" for k in use_keys)
-
-    def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
-        text = text.strip()
-        text_negative = kwargs.get("caption_negative", text).strip()
-        lyrics = kwargs.get("lyrics", "")
-        lyrics_negative = kwargs.get("lyrics_negative", lyrics)
-        duration = kwargs.get("duration", 120)
-        if isinstance(duration, str):
-            duration = float(duration.split(None, 1)[0])
-        language = kwargs.get("language")
-        seed = kwargs.get("seed", 0)
-
-        generate_audio_codes = kwargs.get("generate_audio_codes", True)
-        cfg_scale = kwargs.get("cfg_scale", 2.0)
-        temperature = kwargs.get("temperature", 0.85)
-        top_p = kwargs.get("top_p", 0.9)
-        top_k = kwargs.get("top_k", 0.0)
-        min_p = kwargs.get("min_p", 0.000)
-
-        duration = math.ceil(duration)
-        kwargs["duration"] = duration
-        tokens_duration = duration * 5
-        min_tokens = int(kwargs.get("min_tokens", tokens_duration))
-        max_tokens = int(kwargs.get("max_tokens", tokens_duration))
-
-        metas_negative = {
-            k.rsplit("_", 1)[0]: kwargs.pop(k)
-            for k in ("bpm_negative", "duration_negative", "keyscale_negative", "timesignature_negative", "language_negative", "caption_negative")
-            if k in kwargs
-        }
-        if not kwargs.get("use_negative_caption"):
-            _ = metas_negative.pop("caption", None)
-
-        cot_text = self._metas_to_cot(caption=text, **kwargs)
-        cot_text_negative = "<think>\n\n</think>" if not metas_negative else self._metas_to_cot(**metas_negative)
-        meta_cap = self._metas_to_cap(**kwargs)
-
-        lm_template = "<|im_start|>system\n# Instruction\nGenerate audio semantic tokens based on the given conditions:\n\n<|im_end|>\n<|im_start|>user\n# Caption\n{}\n\n# Lyric\n{}\n<|im_end|>\n<|im_start|>assistant\n{}\n\n<|im_end|>\n"
-        lyrics_template = "# Languages\n{}\n\n# Lyric\n{}<|endoftext|><|endoftext|>"
-        qwen3_06b_template = "# Instruction\nGenerate audio semantic tokens based on the given conditions:\n\n# Caption\n{}\n\n# Metas\n{}\n<|endoftext|>\n<|endoftext|>"
-
-        llm_prompts = {
-            "lm_prompt": lm_template.format(text, lyrics.strip(), cot_text),
-            "lm_prompt_negative": lm_template.format(text_negative, lyrics_negative.strip(), cot_text_negative),
-            "lyrics": lyrics_template.format(language if language is not None else "", lyrics),
-            "qwen3_06b": qwen3_06b_template.format(text, meta_cap),
-        }
-
-        out = {
-            prompt_key: self.qwen3_06b.tokenize_with_weights(
-                prompt,
-                prompt_key == "qwen3_06b" and return_word_ids,
-                disable_weights = True,
-                **kwargs,
-            )
-            for prompt_key, prompt in llm_prompts.items()
-        }
-        out["lm_metadata"] = {"min_tokens": min_tokens,
-                              "max_tokens": max_tokens,
-                              "seed": seed,
-                              "generate_audio_codes": generate_audio_codes,
-                              "cfg_scale": cfg_scale,
-                              "temperature": temperature,
-                              "top_p": top_p,
-                              "top_k": top_k,
-                              "min_p": min_p,
-                              }
-        return out
-
-
-class Qwen3_06BModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen3_06B_ACE15, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-class Qwen3_2B_ACE15(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
-            model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
-
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen3_2B_ACE15_lm, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-class Qwen3_4B_ACE15(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
-            model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
-
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen3_4B_ACE15_lm, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-class ACE15TEModel(torch.nn.Module):
-    def __init__(self, device="cpu", dtype=None, dtype_llama=None, lm_model=None, model_options={}):
-        super().__init__()
-        if dtype_llama is None:
-            dtype_llama = dtype
-
-        model = None
-        self.constant = 0.4375
-        if lm_model == "qwen3_4b":
-            model = Qwen3_4B_ACE15
-            self.constant = 0.5625
-        elif lm_model == "qwen3_2b":
-            model = Qwen3_2B_ACE15
-
-        self.lm_model = lm_model
-        self.qwen3_06b = Qwen3_06BModel(device=device, dtype=dtype, model_options=model_options)
-        if model is not None:
-            setattr(self, self.lm_model, model(device=device, dtype=dtype_llama, model_options=model_options))
-
-        self.dtypes = set([dtype, dtype_llama])
-
-    def encode_token_weights(self, token_weight_pairs):
-        token_weight_pairs_base = token_weight_pairs["qwen3_06b"]
-        token_weight_pairs_lyrics = token_weight_pairs["lyrics"]
-
-        self.qwen3_06b.set_clip_options({"layer": None})
-        base_out, _, extra = self.qwen3_06b.encode_token_weights(token_weight_pairs_base)
-        self.qwen3_06b.set_clip_options({"layer": [0]})
-        lyrics_embeds, _, extra_l = self.qwen3_06b.encode_token_weights(token_weight_pairs_lyrics)
-
-        out = {"conditioning_lyrics": lyrics_embeds[:, 0]}
-
-        lm_metadata = token_weight_pairs["lm_metadata"]
-        if lm_metadata["generate_audio_codes"]:
-            audio_codes = generate_audio_codes(getattr(self, self.lm_model, self.qwen3_06b), token_weight_pairs["lm_prompt"], token_weight_pairs["lm_prompt_negative"], min_tokens=lm_metadata["min_tokens"], max_tokens=lm_metadata["min_tokens"], seed=lm_metadata["seed"], cfg_scale=lm_metadata["cfg_scale"], temperature=lm_metadata["temperature"], top_p=lm_metadata["top_p"], top_k=lm_metadata["top_k"], min_p=lm_metadata["min_p"])
-            out["audio_codes"] = [audio_codes]
-
-        return base_out, None, out
-
-    def set_clip_options(self, options):
-        self.qwen3_06b.set_clip_options(options)
-        lm_model = getattr(self, self.lm_model, None)
-        if lm_model is not None:
-            lm_model.set_clip_options(options)
-
-    def reset_clip_options(self):
-        self.qwen3_06b.reset_clip_options()
-        lm_model = getattr(self, self.lm_model, None)
-        if lm_model is not None:
-            lm_model.reset_clip_options()
-
-    def load_sd(self, sd):
-        if "model.layers.0.post_attention_layernorm.weight" in sd:
-            shape = sd["model.layers.0.post_attention_layernorm.weight"].shape
-            if shape[0] == 1024:
-                return self.qwen3_06b.load_sd(sd)
-            else:
-                return getattr(self, self.lm_model).load_sd(sd)
-
-    def memory_estimation_function(self, token_weight_pairs, device=None):
-        lm_metadata = token_weight_pairs["lm_metadata"]
-        constant = self.constant
-        if comfy.model_management.should_use_bf16(device):
-            constant *= 0.5
-
-        token_weight_pairs = token_weight_pairs.get("lm_prompt", [])
-        num_tokens = sum(map(lambda a: len(a), token_weight_pairs))
-        num_tokens += lm_metadata['min_tokens']
-        return num_tokens * constant * 1024 * 1024
-
-def te(dtype_llama=None, llama_quantization_metadata=None, lm_model="qwen3_2b"):
-    class ACE15TEModel_(ACE15TEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
-                model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
-            super().__init__(device=device, dtype_llama=dtype_llama, lm_model=lm_model, dtype=dtype, model_options=model_options)
-    return ACE15TEModel_
@@ -8,7 +8,7 @@ import torch
 class Qwen3Tokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=1024, embedding_key='qwen3_06b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=151643, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=1024, embedding_key='qwen3_06b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=151643, tokenizer_data=tokenizer_data)

 class T5XXLTokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
@@ -23,7 +23,7 @@ class AnimaTokenizer:
    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs):
        out = {}
        qwen_ids = self.qwen3_06b.tokenize_with_weights(text, return_word_ids, **kwargs)
-        out["qwen3_06b"] = [[(k[0], 1.0, k[2]) if return_word_ids else (k[0], 1.0) for k in inner_list] for inner_list in qwen_ids]  # Set weights to 1.0
+        out["qwen3_06b"] = [[(token, 1.0) for token, _ in inner_list] for inner_list in qwen_ids]  # Set weights to 1.0
        out["t5xxl"] = self.t5xxl.tokenize_with_weights(text, return_word_ids, **kwargs)
        return out

@@ -33,8 +33,6 @@ class AnimaTokenizer:
    def state_dict(self):
        return {}

-    def decode(self, token_ids, **kwargs):
-        return self.qwen3_06b.decode(token_ids, **kwargs)

 class Qwen3_06BModel(sd1_clip.SDClipModel):
    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
@@ -118,7 +118,7 @@ class MistralTokenizerClass:
 class Mistral3Tokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        self.tekken_data = tokenizer_data.get("tekken_model", None)
-        super().__init__("", pad_with_end=False, embedding_directory=embedding_directory, embedding_size=5120, embedding_key='mistral3_24b', tokenizer_class=MistralTokenizerClass, has_end_token=False, pad_to_max_length=False, pad_token=11, start_token=1, max_length=99999999, min_length=1, pad_left=True, tokenizer_args=load_mistral_tokenizer(self.tekken_data), tokenizer_data=tokenizer_data)
+        super().__init__("", pad_with_end=False, embedding_size=5120, embedding_key='mistral3_24b', tokenizer_class=MistralTokenizerClass, has_end_token=False, pad_to_max_length=False, pad_token=11, start_token=1, max_length=99999999, min_length=1, pad_left=True, tokenizer_args=load_mistral_tokenizer(self.tekken_data), tokenizer_data=tokenizer_data)

    def state_dict(self):
        return {"tekken_model": self.tekken_data}
@@ -176,12 +176,12 @@ def flux2_te(dtype_llama=None, llama_quantization_metadata=None, pruned=False):
 class Qwen3Tokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=2560, embedding_key='qwen3_4b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=512, pad_token=151643, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2560, embedding_key='qwen3_4b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=512, pad_token=151643, tokenizer_data=tokenizer_data)

 class Qwen3Tokenizer8B(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=4096, embedding_key='qwen3_8b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=512, pad_token=151643, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=4096, embedding_key='qwen3_8b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=512, pad_token=151643, tokenizer_data=tokenizer_data)

 class KleinTokenizer(sd1_clip.SD1Tokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}, name="qwen3_4b"):
@@ -3,12 +3,9 @@ import torch.nn as nn
 from dataclasses import dataclass
 from typing import Optional, Any, Tuple
 import math
-from tqdm import tqdm
-import comfy.utils

 from comfy.ldm.modules.attention import optimized_attention_for_device
 import comfy.model_management
-import comfy.ops
 import comfy.ldm.common_dit
 import comfy.clip_model

@@ -105,79 +102,6 @@ class Qwen3_06BConfig:
    rope_scale = None
    final_norm: bool = True
    lm_head: bool = False
-    stop_tokens = [151643, 151645]
-
-@dataclass
-class Qwen3_06B_ACE15_Config:
-    vocab_size: int = 151669
-    hidden_size: int = 1024
-    intermediate_size: int = 3072
-    num_hidden_layers: int = 28
-    num_attention_heads: int = 16
-    num_key_value_heads: int = 8
-    max_position_embeddings: int = 32768
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 1000000.0
-    transformer_type: str = "llama"
-    head_dim = 128
-    rms_norm_add = False
-    mlp_activation = "silu"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = "gemma3"
-    k_norm = "gemma3"
-    rope_scale = None
-    final_norm: bool = True
-    lm_head: bool = False
-    stop_tokens = [151643, 151645]
-
-@dataclass
-class Qwen3_2B_ACE15_lm_Config:
-    vocab_size: int = 217204
-    hidden_size: int = 2048
-    intermediate_size: int = 6144
-    num_hidden_layers: int = 28
-    num_attention_heads: int = 16
-    num_key_value_heads: int = 8
-    max_position_embeddings: int = 40960
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 1000000.0
-    transformer_type: str = "llama"
-    head_dim = 128
-    rms_norm_add = False
-    mlp_activation = "silu"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = "gemma3"
-    k_norm = "gemma3"
-    rope_scale = None
-    final_norm: bool = True
-    lm_head: bool = False
-    stop_tokens = [151643, 151645]
-
-@dataclass
-class Qwen3_4B_ACE15_lm_Config:
-    vocab_size: int = 217204
-    hidden_size: int = 2560
-    intermediate_size: int = 9728
-    num_hidden_layers: int = 36
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 8
-    max_position_embeddings: int = 40960
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 1000000.0
-    transformer_type: str = "llama"
-    head_dim = 128
-    rms_norm_add = False
-    mlp_activation = "silu"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = "gemma3"
-    k_norm = "gemma3"
-    rope_scale = None
-    final_norm: bool = True
-    lm_head: bool = False
-    stop_tokens = [151643, 151645]

@dataclass
 class Qwen3_4BConfig:
@@ -201,7 +125,6 @@ class Qwen3_4BConfig:
    rope_scale = None
    final_norm: bool = True
    lm_head: bool = False
-    stop_tokens = [151643, 151645]

@dataclass
 class Qwen3_8BConfig:
@@ -225,7 +148,6 @@ class Qwen3_8BConfig:
    rope_scale = None
    final_norm: bool = True
    lm_head: bool = False
-    stop_tokens = [151643, 151645]

@dataclass
 class Ovis25_2BConfig:
@@ -296,7 +218,6 @@ class Gemma2_2B_Config:
    rope_scale = None
    final_norm: bool = True
    lm_head: bool = False
-    stop_tokens = [1]

@dataclass
 class Gemma3_4B_Config:
@@ -321,14 +242,6 @@ class Gemma3_4B_Config:
    rope_scale = [8.0, 1.0]
    final_norm: bool = True
    lm_head: bool = False
-    stop_tokens = [1, 106]
-
-GEMMA3_VISION_CONFIG = {"num_channels": 3, "hidden_act": "gelu_pytorch_tanh", "hidden_size": 1152, "image_size": 896, "intermediate_size": 4304, "model_type": "siglip_vision_model", "num_attention_heads": 16, "num_hidden_layers": 27, "patch_size": 14}
-
-@dataclass
-class Gemma3_4B_Vision_Config(Gemma3_4B_Config):
-    vision_config = GEMMA3_VISION_CONFIG
-    mm_tokens_per_image = 256

@dataclass
 class Gemma3_12B_Config:
@@ -353,9 +266,8 @@ class Gemma3_12B_Config:
    rope_scale = [8.0, 1.0]
    final_norm: bool = True
    lm_head: bool = False
-    vision_config = GEMMA3_VISION_CONFIG
+    vision_config = {"num_channels": 3, "hidden_act": "gelu_pytorch_tanh", "hidden_size": 1152, "image_size": 896, "intermediate_size": 4304, "model_type": "siglip_vision_model", "num_attention_heads": 16, "num_hidden_layers": 27, "patch_size": 14}
    mm_tokens_per_image = 256
-    stop_tokens = [1, 106]

 class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5, add=False, device=None, dtype=None):
@@ -373,6 +285,13 @@ class RMSNorm(nn.Module):



+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
 def precompute_freqs_cis(head_dim, position_ids, theta, rope_scale=None, rope_dims=None, device=None):
    if not isinstance(theta, list):
        theta = [theta]
@@ -401,30 +320,20 @@ def precompute_freqs_cis(head_dim, position_ids, theta, rope_scale=None, rope_di
        else:
            cos = cos.unsqueeze(1)
            sin = sin.unsqueeze(1)
-        sin_split = sin.shape[-1] // 2
-        out.append((cos, sin[..., : sin_split], -sin[..., sin_split :]))
+        out.append((cos, sin))

    if len(out) == 1:
        return out[0]

    return out

+
 def apply_rope(xq, xk, freqs_cis):
    org_dtype = xq.dtype
    cos = freqs_cis[0]
    sin = freqs_cis[1]
-    nsin = freqs_cis[2]
-
-    q_embed = (xq * cos)
-    q_split = q_embed.shape[-1] // 2
-    q_embed[..., : q_split].addcmul_(xq[..., q_split :], nsin)
-    q_embed[..., q_split :].addcmul_(xq[..., : q_split], sin)
-
-    k_embed = (xk * cos)
-    k_split = k_embed.shape[-1] // 2
-    k_embed[..., : k_split].addcmul_(xk[..., k_split :], nsin)
-    k_embed[..., k_split :].addcmul_(xk[..., : k_split], sin)
-
+    q_embed = (xq * cos) + (rotate_half(xq) * sin)
+    k_embed = (xk * cos) + (rotate_half(xk) * sin)
    return q_embed.to(org_dtype), k_embed.to(org_dtype)


@@ -459,10 +368,8 @@ class Attention(nn.Module):
        freqs_cis: Optional[torch.Tensor] = None,
        optimized_attention=None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        sliding_window: Optional[int] = None,
    ):
        batch_size, seq_length, _ = hidden_states.shape
-
        xq = self.q_proj(hidden_states)
        xk = self.k_proj(hidden_states)
        xv = self.v_proj(hidden_states)
@@ -497,11 +404,6 @@ class Attention(nn.Module):
            else:
                present_key_value = (xk, xv, index + num_tokens)

-            if sliding_window is not None and xk.shape[2] > sliding_window:
-                xk = xk[:, :, -sliding_window:]
-                xv = xv[:, :, -sliding_window:]
-                attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None
-
        xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)

@@ -584,12 +486,10 @@ class TransformerBlockGemma2(nn.Module):
        optimized_attention=None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ):
-        sliding_window = None
        if self.transformer_type == 'gemma3':
            if self.sliding_attention:
-                sliding_window = self.sliding_attention
                if x.shape[1] > self.sliding_attention:
-                    sliding_mask = torch.full((x.shape[1], x.shape[1]), torch.finfo(x.dtype).min, device=x.device, dtype=x.dtype)
+                    sliding_mask = torch.full((x.shape[1], x.shape[1]), float("-inf"), device=x.device, dtype=x.dtype)
                    sliding_mask.tril_(diagonal=-self.sliding_attention)
                    if attention_mask is not None:
                        attention_mask = attention_mask + sliding_mask
@@ -608,7 +508,6 @@ class TransformerBlockGemma2(nn.Module):
            freqs_cis=freqs_cis,
            optimized_attention=optimized_attention,
            past_key_value=past_key_value,
-            sliding_window=sliding_window,
        )

        x = self.post_attention_layernorm(x)
@@ -682,10 +581,10 @@ class Llama2_(nn.Module):
        mask = None
        if attention_mask is not None:
            mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, seq_len, attention_mask.shape[-1])
-            mask = mask.masked_fill(mask.to(torch.bool), torch.finfo(x.dtype).min / 4)
+            mask = mask.masked_fill(mask.to(torch.bool), float("-inf"))

        if seq_len > 1:
-            causal_mask = torch.empty(past_len + seq_len, past_len + seq_len, dtype=x.dtype, device=x.device).fill_(torch.finfo(x.dtype).min / 4).triu_(1)
+            causal_mask = torch.empty(past_len + seq_len, past_len + seq_len, dtype=x.dtype, device=x.device).fill_(float("-inf")).triu_(1)
            if mask is not None:
                mask += causal_mask
            else:
@@ -793,122 +692,6 @@ class BaseLlama:
    def forward(self, input_ids, *args, **kwargs):
        return self.model(input_ids, *args, **kwargs)

-class BaseGenerate:
-    def logits(self, x):
-        input = x[:, -1:]
-        if hasattr(self.model, "lm_head"):
-            module = self.model.lm_head
-        else:
-            module = self.model.embed_tokens
-
-        offload_stream = None
-        if module.comfy_cast_weights:
-            weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True)
-        else:
-            weight = self.model.embed_tokens.weight.to(x)
-
-        x = torch.nn.functional.linear(input, weight, None)
-
-        comfy.ops.uncast_bias_weight(module, weight, None, offload_stream)
-        return x
-
-    def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0):
-        device = embeds.device
-        model_config = self.model.config
-
-        if stop_tokens is None:
-            stop_tokens = self.model.config.stop_tokens
-
-        if execution_dtype is None:
-            if comfy.model_management.should_use_bf16(device):
-                execution_dtype = torch.bfloat16
-            else:
-                execution_dtype = torch.float32
-        embeds = embeds.to(execution_dtype)
-
-        if embeds.ndim == 2:
-            embeds = embeds.unsqueeze(0)
-
-        past_key_values = [] #kv_cache init
-        max_cache_len = embeds.shape[1] + max_length
-        for x in range(model_config.num_hidden_layers):
-            past_key_values.append((torch.empty([embeds.shape[0], model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype),
-                                    torch.empty([embeds.shape[0], model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), 0))
-
-        generator = torch.Generator(device=device).manual_seed(seed) if do_sample else None
-
-        generated_token_ids = []
-        pbar = comfy.utils.ProgressBar(max_length)
-
-        # Generation loop
-        for step in tqdm(range(max_length), desc="Generating tokens"):
-            x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values)
-            logits = self.logits(x)[:, -1]
-            next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, initial_tokens + generated_token_ids, generator, do_sample=do_sample)
-            token_id = next_token[0].item()
-            generated_token_ids.append(token_id)
-
-            embeds = self.model.embed_tokens(next_token).to(execution_dtype)
-            pbar.update(1)
-
-            if token_id in stop_tokens:
-                break
-
-        return generated_token_ids
-
-    def sample_token(self, logits, temperature, top_k, top_p, min_p, repetition_penalty, token_history, generator, do_sample=True):
-
-        if not do_sample or temperature == 0.0:
-            return torch.argmax(logits, dim=-1, keepdim=True)
-
-        # Sampling mode
-        if repetition_penalty != 1.0:
-            for i in range(logits.shape[0]):
-                for token_id in set(token_history):
-                    logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
-
-        if temperature != 1.0:
-            logits = logits / temperature
-
-        if top_k > 0:
-            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
-
-        if min_p > 0.0:
-            probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
-            top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
-            min_threshold = min_p * top_probs
-            indices_to_remove = probs_before_filter < min_threshold
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
-
-        if top_p < 1.0:
-            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-            cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
-            sorted_indices_to_remove = cumulative_probs > top_p
-            sorted_indices_to_remove[..., 0] = False
-            indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
-            indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
-            logits[indices_to_remove] = torch.finfo(logits.dtype).min
-
-        probs = torch.nn.functional.softmax(logits, dim=-1)
-
-        return torch.multinomial(probs, num_samples=1, generator=generator)
-
-class BaseQwen3:
-    def logits(self, x):
-        input = x[:, -1:]
-        module = self.model.embed_tokens
-
-        offload_stream = None
-        if module.comfy_cast_weights:
-            weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True)
-        else:
-            weight = self.model.embed_tokens.weight.to(x)
-
-        x = torch.nn.functional.linear(input, weight, None)
-
-        comfy.ops.uncast_bias_weight(module, weight, None, offload_stream)
-        return x

 class Llama2(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
@@ -937,7 +720,7 @@ class Qwen25_3B(BaseLlama, torch.nn.Module):
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-class Qwen3_06B(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
+class Qwen3_06B(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Qwen3_06BConfig(**config_dict)
@@ -946,25 +729,7 @@ class Qwen3_06B(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-class Qwen3_06B_ACE15(BaseLlama, BaseQwen3, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Qwen3_06B_ACE15_Config(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-
-class Qwen3_2B_ACE15_lm(BaseLlama, BaseQwen3, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Qwen3_2B_ACE15_lm_Config(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-
-class Qwen3_4B(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
+class Qwen3_4B(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Qwen3_4BConfig(**config_dict)
@@ -973,16 +738,7 @@ class Qwen3_4B(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-class Qwen3_4B_ACE15_lm(BaseLlama, BaseQwen3, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Qwen3_4B_ACE15_lm_Config(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-
-class Qwen3_8B(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
+class Qwen3_8B(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Qwen3_8BConfig(**config_dict)
@@ -1000,7 +756,7 @@ class Ovis25_2B(BaseLlama, torch.nn.Module):
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-class Qwen25_7BVLI(BaseLlama, BaseGenerate, torch.nn.Module):
+class Qwen25_7BVLI(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Qwen25_7BVLI_Config(**config_dict)
@@ -1010,9 +766,6 @@ class Qwen25_7BVLI(BaseLlama, BaseGenerate, torch.nn.Module):
        self.visual = qwen_vl.Qwen2VLVisionTransformer(hidden_size=1280, output_hidden_size=config.hidden_size, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-        # todo: should this be tied or not?
-        #self.lm_head = operations.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
-
    def preprocess_embed(self, embed, device):
        if embed["type"] == "image":
            image, grid = qwen_vl.process_qwen2vl_images(embed["data"])
@@ -1046,7 +799,7 @@ class Qwen25_7BVLI(BaseLlama, BaseGenerate, torch.nn.Module):

        return super().forward(x, attention_mask=attention_mask, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=final_layer_norm_intermediate, dtype=dtype, position_ids=position_ids)

-class Gemma2_2B(BaseLlama, BaseGenerate, torch.nn.Module):
+class Gemma2_2B(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Gemma2_2B_Config(**config_dict)
@@ -1055,7 +808,7 @@ class Gemma2_2B(BaseLlama, BaseGenerate, torch.nn.Module):
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-class Gemma3_4B(BaseLlama, BaseGenerate, torch.nn.Module):
+class Gemma3_4B(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Gemma3_4B_Config(**config_dict)
@@ -1064,25 +817,7 @@ class Gemma3_4B(BaseLlama, BaseGenerate, torch.nn.Module):
        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
        self.dtype = dtype

-class Gemma3_4B_Vision(BaseLlama, BaseGenerate, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Gemma3_4B_Vision_Config(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-        self.multi_modal_projector = Gemma3MultiModalProjector(config, dtype, device, operations)
-        self.vision_model = comfy.clip_model.CLIPVision(config.vision_config, dtype, device, operations)
-        self.image_size = config.vision_config["image_size"]
-
-    def preprocess_embed(self, embed, device):
-        if embed["type"] == "image":
-            image = comfy.clip_model.clip_preprocess(embed["data"], size=self.image_size, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], crop=True)
-            return self.multi_modal_projector(self.vision_model(image.to(device, dtype=torch.float32))[0]), None
-        return None, None
-
-class Gemma3_12B(BaseLlama, BaseGenerate, torch.nn.Module):
+class Gemma3_12B(BaseLlama, torch.nn.Module):
    def __init__(self, config_dict, dtype, device, operations):
        super().__init__()
        config = Gemma3_12B_Config(**config_dict)
@@ -6,7 +6,6 @@ import comfy.text_encoders.genmo
 from comfy.ldm.lightricks.embeddings_connector import Embeddings1DConnector
 import torch
 import comfy.utils
-import math

 class T5XXLTokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
@@ -23,79 +22,40 @@ def ltxv_te(*args, **kwargs):
    return comfy.text_encoders.genmo.mochi_te(*args, **kwargs)


-class Gemma3_Tokenizer():
-    def state_dict(self):
-        return {"spiece_model": self.tokenizer.serialize_model()}
-
-    def tokenize_with_weights(self, text, return_word_ids=False, image=None, llama_template=None, skip_template=True, **kwargs):
-        self.llama_template = "<start_of_turn>system\nYou are a helpful assistant.<end_of_turn>\n<start_of_turn>user\n{}<end_of_turn>\n<start_of_turn>model\n"
-        self.llama_template_images = "<start_of_turn>system\nYou are a helpful assistant.<end_of_turn>\n<start_of_turn>user\n\n<image_soft_token>{}<end_of_turn>\n\n<start_of_turn>model\n"
-
-        if image is None:
-            images = []
-        else:
-            samples = image.movedim(-1, 1)
-            total = int(896 * 896)
-
-            scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
-            width = round(samples.shape[3] * scale_by)
-            height = round(samples.shape[2] * scale_by)
-
-            s = comfy.utils.common_upscale(samples, width, height, "area", "disabled").movedim(1, -1)
-            images = [s[:, :, :, :3]]
-
-        if text.startswith('<start_of_turn>'):
-            skip_template = True
-
-        if skip_template:
-            llama_text = text
-        else:
-            if llama_template is None:
-                if len(images) > 0:
-                    llama_text = self.llama_template_images.format(text)
-                else:
-                    llama_text = self.llama_template.format(text)
-            else:
-                llama_text = llama_template.format(text)
-
-        text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
-
-        if len(images) > 0:
-            embed_count = 0
-            for r in text_tokens:
-                for i, token in enumerate(r):
-                    if token[0] == 262144 and embed_count < len(images):
-                        r[i] = ({"type": "image", "data": images[embed_count]},) + token[1:]
-                        embed_count += 1
-        return text_tokens
-
-class Gemma3_12BTokenizer(Gemma3_Tokenizer, sd1_clip.SDTokenizer):
+class Gemma3_12BTokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer = tokenizer_data.get("spiece_model", None)
-        special_tokens = {"<image_soft_token>": 262144, "<end_of_turn>": 106}
-        super().__init__(tokenizer, pad_with_end=False, embedding_size=3840, embedding_key='gemma3_12b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=512, pad_left=True, disable_weights=True, tokenizer_args={"add_bos": True, "add_eos": False, "special_tokens": special_tokens}, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer, pad_with_end=False, embedding_size=3840, embedding_key='gemma3_12b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False}, tokenizer_data=tokenizer_data)

+    def state_dict(self):
+        return {"spiece_model": self.tokenizer.serialize_model()}

 class LTXAVGemmaTokenizer(sd1_clip.SD1Tokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma3_12b", tokenizer=Gemma3_12BTokenizer)

-
 class Gemma3_12BModel(sd1_clip.SDClipModel):
    def __init__(self, device="cpu", layer="all", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
        if llama_quantization_metadata is not None:
            model_options = model_options.copy()
            model_options["quantization_metadata"] = llama_quantization_metadata
-        self.dtypes = set()
-        self.dtypes.add(dtype)
+
        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_12B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)

-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed):
-        tokens_only = [[t[0] for t in b] for b in tokens]
-        embeds, _, _, embeds_info = self.process_tokens(tokens_only, self.execution_device)
-        comfy.utils.normalize_image_embeddings(embeds, embeds_info, self.transformer.model.config.hidden_size ** 0.5)
-        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, stop_tokens=[106])  # 106 is <end_of_turn>
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template="{}", image_embeds=None, **kwargs):
+        text = llama_template.format(text)
+        text_tokens = super().tokenize_with_weights(text, return_word_ids)
+        embed_count = 0
+        for k in text_tokens:
+            tt = text_tokens[k]
+            for r in tt:
+                for i in range(len(r)):
+                    if r[i][0] == 262144:
+                        if image_embeds is not None and embed_count < image_embeds.shape[0]:
+                            r[i] = ({"type": "embedding", "data": image_embeds[embed_count], "original_type": "image"},) + r[i][1:]
+                            embed_count += 1
+        return text_tokens

 class LTXAVTEModel(torch.nn.Module):
    def __init__(self, dtype_llama=None, device="cpu", dtype=None, model_options={}):
@@ -137,7 +97,6 @@ class LTXAVTEModel(torch.nn.Module):
        token_weight_pairs = token_weight_pairs["gemma3_12b"]

        out, pooled, extra = self.gemma3_12b.encode_token_weights(token_weight_pairs)
-        out = out[:, :, -torch.sum(extra["attention_mask"]).item():]
        out_device = out.device
        if comfy.model_management.should_use_bf16(self.execution_device):
            out = out.to(device=self.execution_device, dtype=torch.bfloat16)
@@ -152,9 +111,6 @@ class LTXAVTEModel(torch.nn.Module):

        return out.to(out_device), pooled

-    def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed):
-        return self.gemma3_12b.generate(tokens["gemma3_12b"], do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed)
-
    def load_sd(self, sd):
        if "model.layers.47.self_attn.q_norm.weight" in sd:
            return self.gemma3_12b.load_sd(sd)
@@ -169,7 +125,7 @@ class LTXAVTEModel(torch.nn.Module):
            for prefix, component in [("text_embedding_projection.", self.text_embedding_projection), ("video_embeddings_connector.", self.video_embeddings_connector), ("audio_embeddings_connector.", self.audio_embeddings_connector)]:
                component_sd = {k.replace(prefix, ""): v for k, v in sdo.items() if k.startswith(prefix)}
                if component_sd:
-                    missing, unexpected = component.load_state_dict(component_sd, strict=False, assign=getattr(self, "can_assign_sd", False))
+                    missing, unexpected = component.load_state_dict(component_sd, strict=False)
                    missing_all.extend([f"{prefix}{k}" for k in missing])
                    unexpected_all.extend([f"{prefix}{k}" for k in unexpected])

@@ -182,7 +138,6 @@ class LTXAVTEModel(torch.nn.Module):

        token_weight_pairs = token_weight_pairs.get("gemma3_12b", [])
        num_tokens = sum(map(lambda a: len(a), token_weight_pairs))
-        num_tokens = max(num_tokens, 64)
        return num_tokens * constant * 1024 * 1024

 def ltxav_te(dtype_llama=None, llama_quantization_metadata=None):
@@ -195,14 +150,3 @@ def ltxav_te(dtype_llama=None, llama_quantization_metadata=None):
                dtype = dtype_llama
            super().__init__(dtype_llama=dtype_llama, device=device, dtype=dtype, model_options=model_options)
    return LTXAVTEModel_
-
-def gemma3_te(dtype_llama=None, llama_quantization_metadata=None):
-    class Gemma3_12BModel_(Gemma3_12BModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
-                model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
-            if dtype_llama is not None:
-                dtype = dtype_llama
-            super().__init__(device=device, dtype=dtype, model_options=model_options)
-    return Gemma3_12BModel_
@@ -1,23 +1,23 @@
 from comfy import sd1_clip
 from .spiece_tokenizer import SPieceTokenizer
 import comfy.text_encoders.llama
-from comfy.text_encoders.lt import Gemma3_Tokenizer
-import comfy.utils
+

 class Gemma2BTokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer = tokenizer_data.get("spiece_model", None)
-        special_tokens = {"<end_of_turn>": 107}
-        super().__init__(tokenizer, pad_with_end=False, embedding_size=2304, embedding_key='gemma2_2b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False, "special_tokens": special_tokens}, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer, pad_with_end=False, embedding_size=2304, embedding_key='gemma2_2b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False}, tokenizer_data=tokenizer_data)

    def state_dict(self):
        return {"spiece_model": self.tokenizer.serialize_model()}

-class Gemma3_4BTokenizer(Gemma3_Tokenizer, sd1_clip.SDTokenizer):
+class Gemma3_4BTokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer = tokenizer_data.get("spiece_model", None)
-        special_tokens = {"<image_soft_token>": 262144, "<end_of_turn>": 106}
-        super().__init__(tokenizer, pad_with_end=False, embedding_size=2560, embedding_key='gemma3_4b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False, "special_tokens": special_tokens}, disable_weights=True, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer, pad_with_end=False, embedding_size=2560, embedding_key='gemma3_4b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False}, disable_weights=True, tokenizer_data=tokenizer_data)
+
+    def state_dict(self):
+        return {"spiece_model": self.tokenizer.serialize_model()}

 class LuminaTokenizer(sd1_clip.SD1Tokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
@@ -40,20 +40,6 @@ class Gemma3_4BModel(sd1_clip.SDClipModel):

        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)

-class Gemma3_4B_Vision_Model(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
-            model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
-
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B_Vision, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-    def process_tokens(self, tokens, device):
-        embeds, _, _, embeds_info = super().process_tokens(tokens, device)
-        comfy.utils.normalize_image_embeddings(embeds, embeds_info, self.transformer.model.config.hidden_size ** 0.5)
-        return embeds
-
 class LuminaModel(sd1_clip.SD1ClipModel):
    def __init__(self, device="cpu", dtype=None, model_options={}, name="gemma2_2b", clip_model=Gemma2_2BModel):
        super().__init__(device=device, dtype=dtype, name=name, clip_model=clip_model, model_options=model_options)
@@ -64,8 +50,6 @@ def te(dtype_llama=None, llama_quantization_metadata=None, model_type="gemma2_2b
        model = Gemma2_2BModel
    elif model_type == "gemma3_4b":
        model = Gemma3_4BModel
-    elif model_type == "gemma3_4b_vision":
-        model = Gemma3_4B_Vision_Model

    class LuminaTEModel_(LuminaModel):
        def __init__(self, device="cpu", dtype=None, model_options={}):
@@ -6,10 +6,9 @@ class SPieceTokenizer:
    def from_pretrained(path, **kwargs):
        return SPieceTokenizer(path, **kwargs)

-    def __init__(self, tokenizer_path, add_bos=False, add_eos=True, special_tokens=None):
+    def __init__(self, tokenizer_path, add_bos=False, add_eos=True):
        self.add_bos = add_bos
        self.add_eos = add_eos
-        self.special_tokens = special_tokens
        import sentencepiece
        if torch.is_tensor(tokenizer_path):
            tokenizer_path = tokenizer_path.numpy().tobytes()
@@ -28,32 +27,8 @@ class SPieceTokenizer:
        return out

    def __call__(self, string):
-        if self.special_tokens is not None:
-            import re
-            special_tokens_pattern = '|'.join(re.escape(token) for token in self.special_tokens.keys())
-            if special_tokens_pattern and re.search(special_tokens_pattern, string):
-                parts = re.split(f'({special_tokens_pattern})', string)
-                result = []
-                for part in parts:
-                    if not part:
-                        continue
-                    if part in self.special_tokens:
-                        result.append(self.special_tokens[part])
-                    else:
-                        encoded = self.tokenizer.encode(part, add_bos=False, add_eos=False)
-                        result.extend(encoded)
-                return {"input_ids": result}
-
        out = self.tokenizer.encode(string)
        return {"input_ids": out}

-    def decode(self, token_ids, skip_special_tokens=False):
-
-        if skip_special_tokens and self.special_tokens:
-            special_token_ids = set(self.special_tokens.values())
-            token_ids = [tid for tid in token_ids if tid not in special_token_ids]
-
-        return self.tokenizer.decode(token_ids)
-
    def serialize_model(self):
        return torch.ByteTensor(list(self.tokenizer.serialized_model_proto()))
@@ -6,7 +6,7 @@ import os
 class Qwen3Tokenizer(sd1_clip.SDTokenizer):
    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=2560, embedding_key='qwen3_4b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=151643, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2560, embedding_key='qwen3_4b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=151643, tokenizer_data=tokenizer_data)


 class ZImageTokenizer(sd1_clip.SD1Tokenizer):
--- a/Show More
+++ b/Show More
				`@@ -1 +0,0 @@`
				{"revision": 0, "last_node_id": 29, "last_link_id": 0, "nodes": [{"id": 29, "type": "4c9d6ea4-b912-40e5-8766-6793a9758c53", "pos": [1970, -230], "size": [180, 86], "flags": {}, "order": 5, "mode": 0, "inputs": [{"label": "image", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": null}], "outputs": [{"label": "R", "localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": []}, {"label": "G", "localized_name": "IMAGE1", "name": "IMAGE1", "type": "IMAGE", "links": []}, {"label": "B", "localized_name": "IMAGE2", "name": "IMAGE2", "type": "IMAGE", "links": []}, {"label": "A", "localized_name": "IMAGE3", "name": "IMAGE3", "type": "IMAGE", "links": []}], "title": "Image Channels", "properties": {"proxyWidgets": []}, "widgets_values": []}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "4c9d6ea4-b912-40e5-8766-6793a9758c53", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 28, "lastLinkId": 39, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Image Channels", "inputNode": {"id": -10, "bounding": [1820, -185, 120, 60]}, "outputNode": {"id": -20, "bounding": [2460, -215, 120, 120]}, "inputs": [{"id": "3522932b-2d86-4a1f-a02a-cb29f3a9d7fe", "name": "images.image0", "type": "IMAGE", "linkIds": [39], "localized_name": "images.image0", "label": "image", "pos": [1920, -165]}], "outputs": [{"id": "605cb9c3-b065-4d9b-81d2-3ec331889b2b", "name": "IMAGE0", "type": "IMAGE", "linkIds": [26], "localized_name": "IMAGE0", "label": "R", "pos": [2480, -195]}, {"id": "fb44a77e-0522-43e9-9527-82e7465b3596", "name": "IMAGE1", "type": "IMAGE", "linkIds": [27], "localized_name": "IMAGE1", "label": "G", "pos": [2480, -175]}, {"id": "81460ee6-0131-402a-874f-6bf3001fc4ff", "name": "IMAGE2", "type": "IMAGE", "linkIds": [28], "localized_name": "IMAGE2", "label": "B", "pos": [2480, -155]}, {"id": "ae690246-80d4-4951-b1d9-9306d8a77417", "name": "IMAGE3", "type": "IMAGE", "linkIds": [29], "localized_name": "IMAGE3", "label": "A", "pos": [2480, -135]}], "widgets": [], "nodes": [{"id": 23, "type": "GLSLShader", "pos": [2000, -330], "size": [400, 172], "flags": {}, "order": 0, "mode": 0, "inputs": [{"label": "image", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": 39}, {"localized_name": "fragment_shader", "name": "fragment_shader", "type": "STRING", "widget": {"name": "fragment_shader"}, "link": null}, {"localized_name": "size_mode", "name": "size_mode", "type": "COMFY_DYNAMICCOMBO_V3", "widget": {"name": "size_mode"}, "link": null}, {"label": "image1", "localized_name": "images.image1", "name": "images.image1", "shape": 7, "type": "IMAGE", "link": null}], "outputs": [{"label": "R", "localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": [26]}, {"label": "G", "localized_name": "IMAGE1", "name": "IMAGE1", "type": "IMAGE", "links": [27]}, {"label": "B", "localized_name": "IMAGE2", "name": "IMAGE2", "type": "IMAGE", "links": [28]}, {"label": "A", "localized_name": "IMAGE3", "name": "IMAGE3", "type": "IMAGE", "links": [29]}], "properties": {"Node name for S&R": "GLSLShader"}, "widgets_values": ["#version 300 es\nprecision highp float;\n\nuniform sampler2D u_image0;\n\nin vec2 v_texCoord;\nlayout(location = 0) out vec4 fragColor0;\nlayout(location = 1) out vec4 fragColor1;\nlayout(location = 2) out vec4 fragColor2;\nlayout(location = 3) out vec4 fragColor3;\n\nvoid main() {\n vec4 color = texture(u_image0, v_texCoord);\n // Output each channel as grayscale to separate render targets\n fragColor0 = vec4(vec3(color.r), 1.0); // Red channel\n fragColor1 = vec4(vec3(color.g), 1.0); // Green channel\n fragColor2 = vec4(vec3(color.b), 1.0); // Blue channel\n fragColor3 = vec4(vec3(color.a), 1.0); // Alpha channel\n}\n", "from_input"]}], "groups": [], "links": [{"id": 39, "origin_id": -10, "origin_slot": 0, "target_id": 23, "target_slot": 0, "type": "IMAGE"}, {"id": 26, "origin_id": 23, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "IMAGE"}, {"id": 27, "origin_id": 23, "origin_slot": 1, "target_id": -20, "target_slot": 1, "type": "IMAGE"}, {"id": 28, "origin_id": 23, "origin_slot": 2, "target_id": -20, "target_slot": 2, "type": "IMAGE"}, {"id": 29, "origin_id": 23, "origin_slot": 3, "target_id": -20, "target_slot": 3, "type": "IMAGE"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Image Tools/Color adjust"}]}}
				`@@ -1 +0,0 @@`
				{"revision": 0, "last_node_id": 15, "last_link_id": 0, "nodes": [{"id": 15, "type": "24d8bbfd-39d4-4774-bff0-3de40cc7a471", "pos": [-1490, 2040], "size": [400, 260], "flags": {}, "order": 0, "mode": 0, "inputs": [{"name": "prompt", "type": "STRING", "widget": {"name": "prompt"}, "link": null}, {"label": "reference images", "name": "images", "type": "IMAGE", "link": null}], "outputs": [{"name": "STRING", "type": "STRING", "links": null}], "title": "Prompt Enhance", "properties": {"proxyWidgets": [["-1", "prompt"]], "cnr_id": "comfy-core", "ver": "0.14.1"}, "widgets_values": [""]}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "24d8bbfd-39d4-4774-bff0-3de40cc7a471", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 15, "lastLinkId": 14, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Prompt Enhance", "inputNode": {"id": -10, "bounding": [-2170, 2110, 138.876953125, 80]}, "outputNode": {"id": -20, "bounding": [-640, 2110, 120, 60]}, "inputs": [{"id": "aeab7216-00e0-4528-a09b-bba50845c5a6", "name": "prompt", "type": "STRING", "linkIds": [11], "pos": [-2051.123046875, 2130]}, {"id": "7b73fd36-aa31-4771-9066-f6c83879994b", "name": "images", "type": "IMAGE", "linkIds": [14], "label": "reference images", "pos": [-2051.123046875, 2150]}], "outputs": [{"id": "c7b0d930-68a1-48d1-b496-0519e5837064", "name": "STRING", "type": "STRING", "linkIds": [13], "pos": [-620, 2130]}], "widgets": [], "nodes": [{"id": 11, "type": "GeminiNode", "pos": [-1560, 1990], "size": [470, 470], "flags": {}, "order": 0, "mode": 0, "inputs": [{"localized_name": "images", "name": "images", "shape": 7, "type": "IMAGE", "link": 14}, {"localized_name": "audio", "name": "audio", "shape": 7, "type": "AUDIO", "link": null}, {"localized_name": "video", "name": "video", "shape": 7, "type": "VIDEO", "link": null}, {"localized_name": "files", "name": "files", "shape": 7, "type": "GEMINI_INPUT_FILES", "link": null}, {"localized_name": "prompt", "name": "prompt", "type": "STRING", "widget": {"name": "prompt"}, "link": 11}, {"localized_name": "model", "name": "model", "type": "COMBO", "widget": {"name": "model"}, "link": null}, {"localized_name": "seed", "name": "seed", "type": "INT", "widget": {"name": "seed"}, "link": null}, {"localized_name": "system_prompt", "name": "system_prompt", "shape": 7, "type": "STRING", "widget": {"name": "system_prompt"}, "link": null}], "outputs": [{"localized_name": "STRING", "name": "STRING", "type": "STRING", "links": [13]}], "properties": {"cnr_id": "comfy-core", "ver": "0.14.1", "Node name for S&R": "GeminiNode"}, "widgets_values": ["", "gemini-3-pro-preview", 42, "randomize", "You are an expert in prompt writing.\nBased on the input, rewrite the user's input into a detailed prompt.\nincluding camera settings, lighting, composition, and style.\nReturn the prompt only"], "color": "#432", "bgcolor": "#653"}], "groups": [], "links": [{"id": 11, "origin_id": -10, "origin_slot": 0, "target_id": 11, "target_slot": 4, "type": "STRING"}, {"id": 13, "origin_id": 11, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "STRING"}, {"id": 14, "origin_id": -10, "origin_slot": 1, "target_id": 11, "target_slot": 0, "type": "IMAGE"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Text generation/Prompt enhance"}]}, "extra": {}}
				`@@ -1 +0,0 @@`
				{"revision": 0, "last_node_id": 25, "last_link_id": 0, "nodes": [{"id": 25, "type": "621ba4e2-22a8-482d-a369-023753198b7b", "pos": [4610, -790], "size": [230, 58], "flags": {}, "order": 4, "mode": 0, "inputs": [{"label": "image", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": null}], "outputs": [{"label": "IMAGE", "localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": []}], "title": "Sharpen", "properties": {"proxyWidgets": [["24", "value"]]}, "widgets_values": []}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "621ba4e2-22a8-482d-a369-023753198b7b", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 24, "lastLinkId": 36, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Sharpen", "inputNode": {"id": -10, "bounding": [4090, -825, 120, 60]}, "outputNode": {"id": -20, "bounding": [5150, -825, 120, 60]}, "inputs": [{"id": "37011fb7-14b7-4e0e-b1a0-6a02e8da1fd7", "name": "images.image0", "type": "IMAGE", "linkIds": [34], "localized_name": "images.image0", "label": "image", "pos": [4190, -805]}], "outputs": [{"id": "e9182b3f-635c-4cd4-a152-4b4be17ae4b9", "name": "IMAGE0", "type": "IMAGE", "linkIds": [35], "localized_name": "IMAGE0", "label": "IMAGE", "pos": [5170, -805]}], "widgets": [], "nodes": [{"id": 24, "type": "PrimitiveFloat", "pos": [4280, -1240], "size": [270, 58], "flags": {}, "order": 0, "mode": 0, "inputs": [{"label": "strength", "localized_name": "value", "name": "value", "type": "FLOAT", "widget": {"name": "value"}, "link": null}], "outputs": [{"localized_name": "FLOAT", "name": "FLOAT", "type": "FLOAT", "links": [36]}], "properties": {"Node name for S&R": "PrimitiveFloat", "min": 0, "max": 3, "precision": 2, "step": 0.05}, "widgets_values": [0.5]}, {"id": 23, "type": "GLSLShader", "pos": [4570, -1240], "size": [370, 192], "flags": {}, "order": 1, "mode": 0, "inputs": [{"label": "image0", "localized_name": "images.image0", "name": "images.image0", "type": "IMAGE", "link": 34}, {"label": "image1", "localized_name": "images.image1", "name": "images.image1", "shape": 7, "type": "IMAGE", "link": null}, {"label": "u_float0", "localized_name": "floats.u_float0", "name": "floats.u_float0", "shape": 7, "type": "FLOAT", "link": 36}, {"label": "u_float1", "localized_name": "floats.u_float1", "name": "floats.u_float1", "shape": 7, "type": "FLOAT", "link": null}, {"label": "u_int0", "localized_name": "ints.u_int0", "name": "ints.u_int0", "shape": 7, "type": "INT", "link": null}, {"localized_name": "fragment_shader", "name": "fragment_shader", "type": "STRING", "widget": {"name": "fragment_shader"}, "link": null}, {"localized_name": "size_mode", "name": "size_mode", "type": "COMFY_DYNAMICCOMBO_V3", "widget": {"name": "size_mode"}, "link": null}], "outputs": [{"localized_name": "IMAGE0", "name": "IMAGE0", "type": "IMAGE", "links": [35]}, {"localized_name": "IMAGE1", "name": "IMAGE1", "type": "IMAGE", "links": null}, {"localized_name": "IMAGE2", "name": "IMAGE2", "type": "IMAGE", "links": null}, {"localized_name": "IMAGE3", "name": "IMAGE3", "type": "IMAGE", "links": null}], "properties": {"Node name for S&R": "GLSLShader"}, "widgets_values": ["#version 300 es\nprecision highp float;\n\nuniform sampler2D u_image0;\nuniform vec2 u_resolution;\nuniform float u_float0; // strength [0.0 – 2.0] typical: 0.3–1.0\n\nin vec2 v_texCoord;\nlayout(location = 0) out vec4 fragColor0;\n\nvoid main() {\n vec2 texel = 1.0 / u_resolution;\n \n // Sample center and neighbors\n vec4 center = texture(u_image0, v_texCoord);\n vec4 top = texture(u_image0, v_texCoord + vec2( 0.0, -texel.y));\n vec4 bottom = texture(u_image0, v_texCoord + vec2( 0.0, texel.y));\n vec4 left = texture(u_image0, v_texCoord + vec2(-texel.x, 0.0));\n vec4 right = texture(u_image0, v_texCoord + vec2( texel.x, 0.0));\n \n // Edge enhancement (Laplacian)\n vec4 edges = center * 4.0 - top - bottom - left - right;\n \n // Add edges back scaled by strength\n vec4 sharpened = center + edges * u_float0;\n \n fragColor0 = vec4(clamp(sharpened.rgb, 0.0, 1.0), center.a);\n}", "from_input"]}], "groups": [], "links": [{"id": 36, "origin_id": 24, "origin_slot": 0, "target_id": 23, "target_slot": 2, "type": "FLOAT"}, {"id": 34, "origin_id": -10, "origin_slot": 0, "target_id": 23, "target_slot": 0, "type": "IMAGE"}, {"id": 35, "origin_id": 23, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "IMAGE"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Image Tools/Sharpen"}]}}
				`@@ -1 +0,0 @@`
				{"revision": 0, "last_node_id": 13, "last_link_id": 0, "nodes": [{"id": 13, "type": "cf95b747-3e17-46cb-8097-cac60ff9b2e1", "pos": [1120, 330], "size": [240, 58], "flags": {}, "order": 3, "mode": 0, "inputs": [{"localized_name": "video", "name": "video", "type": "VIDEO", "link": null}, {"name": "model_name", "type": "COMBO", "widget": {"name": "model_name"}, "link": null}], "outputs": [{"localized_name": "VIDEO", "name": "VIDEO", "type": "VIDEO", "links": []}], "title": "Video Upscale(GAN x4)", "properties": {"proxyWidgets": [["-1", "model_name"]], "cnr_id": "comfy-core", "ver": "0.14.1"}, "widgets_values": ["RealESRGAN_x4plus.safetensors"]}], "links": [], "version": 0.4, "definitions": {"subgraphs": [{"id": "cf95b747-3e17-46cb-8097-cac60ff9b2e1", "version": 1, "state": {"lastGroupId": 0, "lastNodeId": 13, "lastLinkId": 19, "lastRerouteId": 0}, "revision": 0, "config": {}, "name": "Video Upscale(GAN x4)", "inputNode": {"id": -10, "bounding": [550, 460, 120, 80]}, "outputNode": {"id": -20, "bounding": [1490, 460, 120, 60]}, "inputs": [{"id": "666d633e-93e7-42dc-8d11-2b7b99b0f2a6", "name": "video", "type": "VIDEO", "linkIds": [10], "localized_name": "video", "pos": [650, 480]}, {"id": "2e23a087-caa8-4d65-99e6-662761aa905a", "name": "model_name", "type": "COMBO", "linkIds": [19], "pos": [650, 500]}], "outputs": [{"id": "0c1768ea-3ec2-412f-9af6-8e0fa36dae70", "name": "VIDEO", "type": "VIDEO", "linkIds": [15], "localized_name": "VIDEO", "pos": [1510, 480]}], "widgets": [], "nodes": [{"id": 2, "type": "ImageUpscaleWithModel", "pos": [1110, 450], "size": [320, 46], "flags": {}, "order": 1, "mode": 0, "inputs": [{"localized_name": "upscale_model", "name": "upscale_model", "type": "UPSCALE_MODEL", "link": 1}, {"localized_name": "image", "name": "image", "type": "IMAGE", "link": 14}], "outputs": [{"localized_name": "IMAGE", "name": "IMAGE", "type": "IMAGE", "links": [13]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "ImageUpscaleWithModel"}}, {"id": 11, "type": "CreateVideo", "pos": [1110, 550], "size": [320, 78], "flags": {}, "order": 3, "mode": 0, "inputs": [{"localized_name": "images", "name": "images", "type": "IMAGE", "link": 13}, {"localized_name": "audio", "name": "audio", "shape": 7, "type": "AUDIO", "link": 16}, {"localized_name": "fps", "name": "fps", "type": "FLOAT", "widget": {"name": "fps"}, "link": 12}], "outputs": [{"localized_name": "VIDEO", "name": "VIDEO", "type": "VIDEO", "links": [15]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "CreateVideo"}, "widgets_values": [30]}, {"id": 10, "type": "GetVideoComponents", "pos": [1110, 330], "size": [320, 70], "flags": {}, "order": 2, "mode": 0, "inputs": [{"localized_name": "video", "name": "video", "type": "VIDEO", "link": 10}], "outputs": [{"localized_name": "images", "name": "images", "type": "IMAGE", "links": [14]}, {"localized_name": "audio", "name": "audio", "type": "AUDIO", "links": [16]}, {"localized_name": "fps", "name": "fps", "type": "FLOAT", "links": [12]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "GetVideoComponents"}}, {"id": 1, "type": "UpscaleModelLoader", "pos": [750, 450], "size": [280, 60], "flags": {}, "order": 0, "mode": 0, "inputs": [{"localized_name": "model_name", "name": "model_name", "type": "COMBO", "widget": {"name": "model_name"}, "link": 19}], "outputs": [{"localized_name": "UPSCALE_MODEL", "name": "UPSCALE_MODEL", "type": "UPSCALE_MODEL", "links": [1]}], "properties": {"cnr_id": "comfy-core", "ver": "0.10.0", "Node name for S&R": "UpscaleModelLoader", "models": [{"name": "RealESRGAN_x4plus.safetensors", "url": "https://huggingface.co/Comfy-Org/Real-ESRGAN_repackaged/resolve/main/RealESRGAN_x4plus.safetensors", "directory": "upscale_models"}]}, "widgets_values": ["RealESRGAN_x4plus.safetensors"]}], "groups": [], "links": [{"id": 1, "origin_id": 1, "origin_slot": 0, "target_id": 2, "target_slot": 0, "type": "UPSCALE_MODEL"}, {"id": 14, "origin_id": 10, "origin_slot": 0, "target_id": 2, "target_slot": 1, "type": "IMAGE"}, {"id": 13, "origin_id": 2, "origin_slot": 0, "target_id": 11, "target_slot": 0, "type": "IMAGE"}, {"id": 16, "origin_id": 10, "origin_slot": 1, "target_id": 11, "target_slot": 1, "type": "AUDIO"}, {"id": 12, "origin_id": 10, "origin_slot": 2, "target_id": 11, "target_slot": 2, "type": "FLOAT"}, {"id": 10, "origin_id": -10, "origin_slot": 0, "target_id": 10, "target_slot": 0, "type": "VIDEO"}, {"id": 15, "origin_id": 11, "origin_slot": 0, "target_id": -20, "target_slot": 0, "type": "VIDEO"}, {"id": 19, "origin_id": -10, "origin_slot": 1, "target_id": 1, "target_slot": 0, "type": "COMBO"}], "extra": {"workflowRendererVersion": "LG"}, "category": "Video generation and editing/Enhance video"}]}, "extra": {}}