Move NF4 Gemma encoder into app/, patch ltx_pipelines at runtime

app/service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import functools
 import logging
 import os
 import uuid
@@ -12,14 +13,68 @@ import torch
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_pipelines import DistilledPipeline
 from ltx_pipelines.utils.args import ImageConditioningInput
+from ltx_pipelines.utils.blocks import PromptEncoder as UpstreamPromptEncoder
 from PIL import Image
 
 from app.config import settings
 from app.database import JobDB, JobRecord
+from app.text_encoder import Nf4PromptEncoder
 
 logger = logging.getLogger(__name__)
 
 
+def _patch_ltx_pipelines() -> None:
+    """Inject our Nf4PromptEncoder and wire DistilledPipeline to accept gemma_quantization."""
+    import types
+
+    import ltx_pipelines.distilled as _distilled_mod
+    import ltx_pipelines.utils.blocks as _blocks_mod
+
+    _blocks_mod.PromptEncoder = Nf4PromptEncoder  # type: ignore[attr-defined]
+
+    original_init = DistilledPipeline.__init__
+
+    @functools.wraps(original_init)
+    def patched_init(
+        self,
+        distilled_checkpoint_path,
+        gemma_root,
+        spatial_upsampler_path,
+        loras,
+        device=None,
+        quantization=None,
+        registry=None,
+        compilation_config=None,
+        offload_mode=None,
+        gemma_quantization=None,
+    ):
+        original_init(
+            self,
+            distilled_checkpoint_path=distilled_checkpoint_path,
+            gemma_root=gemma_root,
+            spatial_upsampler_path=spatial_upsampler_path,
+            loras=loras,
+            device=device,
+            quantization=quantization,
+            registry=registry,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
+        )
+        # Rebuild prompt_encoder with gemma_quantization
+        self.prompt_encoder = Nf4PromptEncoder(
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
+            gemma_quantization=gemma_quantization,
+        )
+
+    DistilledPipeline.__init__ = patched_init  # type: ignore[method-assign]
+    _distilled_mod.DistilledPipeline.__init__ = patched_init  # type: ignore[method-assign]
+
+
 class LtxService:
     def __init__(self, db: JobDB) -> None:
         self.db = db
@@ -32,6 +87,7 @@ class LtxService:
         async with self._load_lock:
             if self.pipeline is not None:
                 return
+            _patch_ltx_pipelines()
             logger.info("Loading DistilledPipeline...")
             lora_config = self._parse_lora_paths(settings.ltx_loras)
             quantization = self._build_quantization_policy()