diff --git a/app/service.py b/app/service.py index e90e7b4..5aca466 100644 --- a/app/service.py +++ b/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() diff --git a/app/text_encoder.py b/app/text_encoder.py new file mode 100644 index 0000000..5ad3526 --- /dev/null +++ b/app/text_encoder.py @@ -0,0 +1,210 @@ +"""NF4 quantized Gemma prompt encoder – owned by this project, patches ltx_pipelines at runtime.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager + +import torch + +from ltx_core.block_streaming import StreamingModelBuilder +from ltx_core.loader.primitives import BuilderProtocol +from ltx_core.loader.registry import DummyRegistry, Registry +from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder +from ltx_core.text_encoders.gemma import ( + EMBEDDINGS_PROCESSOR_KEY_OPS, + GEMMA_LLM_KEY_OPS, + GEMMA_MODEL_OPS, + EmbeddingsProcessorConfigurator, + GemmaTextEncoderConfigurator, + GemmaTextEncoder, + module_ops_from_gemma_root, +) +from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor, EmbeddingsProcessorOutput +from ltx_core.text_encoders.gemma.tokenizer import GemmaTokenizer as LTXVGemmaTokenizer +from ltx_core.utils import find_matching_file +from ltx_pipelines.utils.blocks import PromptEncoder as BasePromptEncoder, _streaming_model +from ltx_pipelines.utils.gpu_model import gpu_model +from ltx_pipelines.utils.helpers import generate_enhanced_prompt +from ltx_pipelines.utils.types import OffloadMode + +logger = logging.getLogger(__name__) + + +class Nf4PromptEncoder(BasePromptEncoder): + """PromptEncoder with NF4 quantized Gemma support.""" + + def __init__( + self, + checkpoint_path: str, + gemma_root: str, + dtype: torch.dtype, + device: torch.device, + registry: Registry | None = None, + offload_mode: OffloadMode = OffloadMode.NONE, + text_encoder_builder: BuilderProtocol | None = None, + gemma_quantization: str | None = None, + ) -> None: + self._gemma_root = gemma_root + self._checkpoint_path = checkpoint_path + self._dtype = dtype + self._device = device + self._offload_mode = offload_mode + self._gemma_quantization = gemma_quantization + self._cached_quantized_encoder: GemmaTextEncoder | None = None + self._processor = None # type: ignore + + if gemma_quantization == "nf4": + self._is_quantized = True + elif gemma_quantization is not None and gemma_quantization != "": + logger.warning( + "Unsupported gemma_quantization '%s', falling back to bf16", gemma_quantization + ) + self._is_quantized = False + else: + self._is_quantized = False + + if text_encoder_builder is not None: + if offload_mode != OffloadMode.NONE: + raise ValueError( + "text_encoder_builder cannot be used with offload_mode != OffloadMode.NONE" + ) + self._text_encoder_builder = text_encoder_builder + self._streaming_text_encoder_builder = None + elif not self._is_quantized: + module_ops = module_ops_from_gemma_root(gemma_root) + model_folder = find_matching_file(gemma_root, "model*.safetensors").parent + weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")] + self._text_encoder_builder = Builder( + model_path=tuple(weight_paths), + model_class_configurator=GemmaTextEncoderConfigurator, + model_sd_ops=GEMMA_LLM_KEY_OPS, + module_ops=(GEMMA_MODEL_OPS, *module_ops), + registry=registry or DummyRegistry(), + ) + self._streaming_text_encoder_builder = StreamingModelBuilder( + model_path=tuple(weight_paths), + model_class_configurator=GemmaTextEncoderConfigurator, + model_sd_ops=GEMMA_LLM_KEY_OPS, + module_ops=(GEMMA_MODEL_OPS, *module_ops), + registry=registry or DummyRegistry(), + blocks_attr="model.model.language_model.layers", + blocks_prefix="model.model.language_model.layers", + ) + + self._embeddings_processor_builder = Builder( + model_path=checkpoint_path, + model_class_configurator=EmbeddingsProcessorConfigurator, + model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS, + registry=registry or DummyRegistry(), + ) + + @staticmethod + def _try_import_bitsandbytes() -> None: + try: + import bitsandbytes # noqa: F401 + except ImportError as e: + raise RuntimeError( + "NF4 Gemma quantization requires bitsandbytes. Install with:\n" + " pip install bitsandbytes\n" + "Or disable REVIDS_LTX_GEMMA_QUANTIZATION to use bf16." + ) from e + + def _load_quantized_gemma(self) -> GemmaTextEncoder: + self._try_import_bitsandbytes() + from transformers import ( + AutoImageProcessor, + BitsAndBytesConfig, + Gemma3ForConditionalGeneration, + Gemma3Processor, + ) + + gemma_path = str(find_matching_file(self._gemma_root, "model*.safetensors").parent) + tokenizer_path = str(find_matching_file(self._gemma_root, "tokenizer.model").parent) + processor_path = str(find_matching_file(self._gemma_root, "preprocessor_config.json").parent) + + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + llm_int8_skip_modules=["lm_head"], + ) + + with torch.device("meta"): + gemma_model = Gemma3ForConditionalGeneration.from_pretrained( + gemma_path, + quantization_config=bnb_config, + torch_dtype=torch.bfloat16, + device_map={"": self._device}, + local_files_only=True, + ) + + tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024) + image_processor = AutoImageProcessor.from_pretrained( + processor_path, local_files_only=True, use_fast=False + ) + self._processor = Gemma3Processor( + image_processor=image_processor, tokenizer=tokenizer.tokenizer + ) + + return GemmaTextEncoder( + model=gemma_model, tokenizer=tokenizer, processor=self._processor, dtype=self._dtype + ) + + def _build_text_encoder(self) -> torch.nn.Module: + if self._is_quantized: + if self._cached_quantized_encoder is None: + logger.info("Loading NF4 quantized Gemma encoder from %s", self._gemma_root) + self._cached_quantized_encoder = self._load_quantized_gemma() + logger.info("NF4 Gemma encoder loaded and cached") + return self._cached_quantized_encoder.eval() + return self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval() + + def _build_embeddings_processor(self) -> EmbeddingsProcessor: + return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).eval() + + def _text_encoder_ctx(self) -> AbstractContextManager: + if self._is_quantized: + + @contextmanager + def _cached_ctx() -> Iterator: + yield self._build_text_encoder() + + return _cached_ctx() + if self._offload_mode != OffloadMode.NONE: + return _streaming_model( + self._streaming_text_encoder_builder, + self._offload_mode, + self._device, + self._dtype, + ) + return gpu_model(self._build_text_encoder()) + + def __call__( + self, + prompts: list[str], + *, + enhance_first_prompt: bool = False, + enhance_prompt_image: str | None = None, + enhance_prompt_seed: int = 42, + ) -> list[EmbeddingsProcessorOutput]: + logger.info("Building text encoder from %s", self._gemma_root) + with self._text_encoder_ctx() as text_encoder: # type: ignore[var-annotated] + if enhance_first_prompt: + prompts = list(prompts) + prompts[0] = generate_enhanced_prompt( + text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed + ) + raw_outputs = [text_encoder.encode(p) for p in prompts] + logger.info( + "Text encoder done, building embeddings processor from %s", self._checkpoint_path + ) + + with gpu_model(self._build_embeddings_processor()) as embeddings_processor: + result = [ + embeddings_processor.process_hidden_states(hs, mask) + for hs, mask in raw_outputs + ] + logger.info("Prompt encoding complete") + return result