Load raw Gemma weights with Sulphur key mapping

app.py

@@ -86,9 +86,11 @@ builtins.__import__ = _import_without_xformers
 try:
     from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
     from ltx_core.loader.primitives import StateDict
+    from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
     from ltx_core.loader.sft_loader import SafetensorsStateDictLoader
     from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
     from ltx_core.quantization import QuantizationPolicy
+    import ltx_pipelines.utils.model_ledger as _model_ledger_mod
     from ltx_pipelines.distilled import DistilledPipeline
     from ltx_pipelines.utils.args import ImageConditioningInput
     from ltx_pipelines.utils.media_io import encode_video
@@ -96,6 +98,24 @@ finally:
     builtins.__import__ = _real_import
 
 
+_model_ledger_mod.GEMMA_LLM_KEY_OPS = (
+    SDOps("GEMMA_LLM_KEY_OPS_SULPHUR")
+    .with_matching(prefix="model.language_model.")
+    .with_matching(prefix="model.vision_tower.")
+    .with_matching(prefix="model.multi_modal_projector.")
+    .with_replacement("model.language_model.", "model.model.language_model.")
+    .with_replacement("model.vision_tower.", "model.model.vision_tower.")
+    .with_replacement("model.multi_modal_projector.", "model.model.multi_modal_projector.")
+    .with_kv_operation(
+        operation=lambda key, value: [
+            KeyValueOperationResult(key, value),
+            KeyValueOperationResult("model.lm_head.weight", value),
+        ],
+        key_prefix="model.model.language_model.embed_tokens.weight",
+    )
+)
+
+
 _SAFETENSORS_DTYPE_MAP = {
     "F64": torch.float64,
     "F32": torch.float32,
@@ -168,12 +188,11 @@ def _download(repo_id, filename, local_dir):
 def _ensure_gemma_model_alias(gemma_weight):
     gemma_weight = Path(gemma_weight)
     alias = gemma_weight.with_name("model.safetensors")
-    if alias.exists():
+    if gemma_weight == alias:
         return alias
-    try:
-        alias.symlink_to(gemma_weight.name)
-    except OSError:
-        os.link(gemma_weight, alias)
+    if alias.exists() or alias.is_symlink():
+        alias.unlink()
+    gemma_weight.replace(alias)
     print(f"[download] ready: {alias}")
     return alias
 
@@ -198,9 +217,10 @@ def download_assets():
 
     gemma_folder = "gemma-3-12b-it-qat-q4_0-unquantized"
     gemma_files = [
-        f"{gemma_folder}/{gemma_folder}_quanto_bf16_int8.safetensors",
+        f"{gemma_folder}/{gemma_folder}.safetensors",
         f"{gemma_folder}/added_tokens.json",
         f"{gemma_folder}/chat_template.json",
+        f"{gemma_folder}/config.json",
         f"{gemma_folder}/config_light.json",
         f"{gemma_folder}/generation_config.json",
         f"{gemma_folder}/preprocessor_config.json",
@@ -216,6 +236,9 @@ def download_assets():
         if filename.endswith(".safetensors"):
             gemma_weight = path
 
+    old_quanto = ASSETS_DIR / gemma_folder / f"{gemma_folder}_quanto_bf16_int8.safetensors"
+    if old_quanto.exists():
+        old_quanto.unlink()
     _ensure_gemma_model_alias(gemma_weight)
 
     return {