Update app.py

app.py

@@ -295,6 +295,73 @@ checkpoint_path = hf_hub_download(
     local_dir=str(weights_dir),
     local_dir_use_symlinks=False,
 )
+
+print("[Gemma] Setting up abliterated Gemma text encoder...")
+MERGED_WEIGHTS = "/tmp/abliterated_gemma_merged.safetensors"
+gemma_root = "/tmp/abliterated_gemma"
+os.makedirs(gemma_root, exist_ok=True)
+
+gemma_official_dir = snapshot_download(
+    repo_id=GEMMA_REPO,
+    ignore_patterns=["*.safetensors", "*.safetensors.index.json"],
+)
+
+for fname in os.listdir(gemma_official_dir):
+    src = os.path.join(gemma_official_dir, fname)
+    dst = os.path.join(gemma_root, fname)
+    if os.path.isfile(src) and not fname.endswith(".safetensors") and fname != "model.safetensors.index.json":
+        if not os.path.exists(dst):
+            os.symlink(src, dst)
+
+if os.path.exists(MERGED_WEIGHTS):
+    print("[Gemma] Using cached merged weights")
+else:
+    abliterated_weights_path = hf_hub_download(
+        repo_id=GEMMA_ABLITERATED_REPO,
+        filename=GEMMA_ABLITERATED_FILE,
+    )
+    index_path = hf_hub_download(
+        repo_id=GEMMA_REPO,
+        filename="model.safetensors.index.json"
+    )
+    with open(index_path) as f:
+        weight_index = json.load(f)
+
+    vision_keys = {}
+    for key, shard in weight_index["weight_map"].items():
+        if "vision_tower" in key or "multi_modal_projector" in key:
+            vision_keys[key] = shard
+    needed_shards = set(vision_keys.values())
+
+    shard_paths = {}
+    for shard_name in needed_shards:
+        shard_paths[shard_name] = hf_hub_download(
+            repo_id=GEMMA_REPO,
+            filename=shard_name
+        )
+
+    _fp8_types = {torch.float8_e4m3fn, torch.float8_e5m2}
+    raw = load_file(abliterated_weights_path)
+    merged = {}
+    for key, tensor in raw.items():
+        t = tensor.to(torch.bfloat16) if tensor.dtype in _fp8_types else tensor
+        merged[f"language_model.{key}"] = t
+    del raw
+
+    for key, shard_name in vision_keys.items():
+        with safe_open(shard_paths[shard_name], framework="pt") as f:
+            merged[key] = f.get_tensor(key)
+
+    save_file(merged, MERGED_WEIGHTS)
+    del merged
+    gc.collect()
+
+weight_link = os.path.join(gemma_root, "model.safetensors")
+if os.path.exists(weight_link):
+    os.remove(weight_link)
+os.symlink(MERGED_WEIGHTS, weight_link)
+print(f"[Gemma] Root ready: {gemma_root}")
+
 spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
 gemma_root = snapshot_download(repo_id=GEMMA_REPO)