Faz2: code/meshai_train/base_weights.py

code/meshai_train/base_weights.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+"""HF MeshAI-Base-Models (dataset) uzerinden TRELLIS/Hunyuan agirlik indirme + LoRA hedef secimi."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+BASE_DATASET_REPO = "HayrettinIscan/MeshAI-Base-Models"
+
+# A100 icin makul ilk paket (Hunyuan DiT ~5GB — VM disk yeterli)
+DEFAULT_WEIGHT_FILES = (
+    "Microsoft-TRELLIS/ckpts/slat_enc_swin8_B_64l8_fp16.safetensors",
+    "Microsoft-TRELLIS/ckpts/ss_enc_conv3d_16l8_fp16.safetensors",
+    "Microsoft-TRELLIS/ckpts/slat_dec_mesh_swin8_B_64l8m256c_fp16.safetensors",
+    "Tencent-Hunyuan3D/hunyuan3d-dit-v2-0-turbo/model.fp16.safetensors",
+)
+
+
+def ensure_base_weights(
+    *,
+    token: str,
+    cache_dir: Path,
+    files: tuple[str, ...] = DEFAULT_WEIGHT_FILES,
+    log_fn: Any = print,
+) -> dict[str, Path]:
+    """Indirilen safetensors yollarini dondurur (HF hub cache veya local_dir)."""
+    from huggingface_hub import hf_hub_download
+
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    out: dict[str, Path] = {}
+    for rel in files:
+        log_fn(f"[faz2] indiriliyor: {rel}")
+        path = Path(
+            hf_hub_download(
+                repo_id=BASE_DATASET_REPO,
+                filename=rel,
+                repo_type="dataset",
+                token=token,
+                local_dir=str(cache_dir),
+            )
+        )
+        # hf bazen local_dir/rel yazar
+        candidate = cache_dir / rel
+        if candidate.exists():
+            path = candidate
+        out[rel] = path
+        log_fn(f"[faz2] hazir: {path} ({path.stat().st_size // (1024 * 1024)} MB)")
+    meta = cache_dir / "faz2_weight_manifest.json"
+    meta.write_text(
+        json.dumps({k: str(v) for k, v in out.items()}, indent=2),
+        encoding="utf-8",
+    )
+    return out
+
+
+def _load_safetensors(path: Path) -> dict[str, torch.Tensor]:
+    try:
+        from safetensors.torch import load_file
+
+        return load_file(str(path))
+    except Exception:
+        # fallback: torch.load for .ckpt/.pt
+        return torch.load(path, map_location="cpu", weights_only=True)
+
+
+def pick_lora_targets(
+    state: dict[str, torch.Tensor],
+    *,
+    max_matrices: int = 8,
+    min_in: int = 64,
+    max_in: int = 8192,
+) -> list[tuple[str, torch.Tensor]]:
+    """2D weight matrislerinden LoRA adaylarini sec (buyukten kucuge)."""
+    cands: list[tuple[str, torch.Tensor, int]] = []
+    for key, tensor in state.items():
+        if not key.endswith("weight"):
+            continue
+        if tensor.ndim != 2:
+            continue
+        out_f, in_f = int(tensor.shape[0]), int(tensor.shape[1])
+        if in_f < min_in or in_f > max_in:
+            continue
+        if out_f < 16:
+            continue
+        cands.append((key, tensor.detach().float().cpu(), out_f * in_f))
+    cands.sort(key=lambda x: x[2], reverse=True)
+    return [(k, t) for k, t, _ in cands[:max_matrices]]
+
+
+class FrozenLinearLoRA(nn.Module):
+    """W frozen + LoRA(A,B). y = x @ W.T + scale * (x @ A.T) @ B.T"""
+
+    def __init__(
+        self,
+        weight: torch.Tensor,
+        *,
+        rank: int = 8,
+        scale: float = 1.0,
+        name: str = "",
+    ) -> None:
+        super().__init__()
+        out_f, in_f = weight.shape
+        self.name = name
+        self.in_features = in_f
+        self.out_features = out_f
+        self.scale = scale
+        self.register_buffer("weight", weight.contiguous())
+        self.lora_A = nn.Parameter(torch.zeros(rank, in_f))
+        self.lora_B = nn.Parameter(torch.zeros(out_f, rank))
+        nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
+        nn.init.zeros_(self.lora_B)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        base = nn.functional.linear(x, self.weight)
+        delta = (x @ self.lora_A.T) @ self.lora_B.T
+        return base + self.scale * delta
+
+
+class BaseLoRATower(nn.Module):
+    """Birden fazla frozen+LoRA katmani; giris projesi ile dim hizalama."""
+
+    def __init__(
+        self,
+        targets: list[tuple[str, torch.Tensor]],
+        *,
+        input_dim: int,
+        rank: int = 8,
+        out_dim: int = 512,
+    ) -> None:
+        super().__init__()
+        if not targets:
+            raise ValueError("LoRA hedefi yok")
+        layers = nn.ModuleList()
+        # Ilk katman in_features'a proj_in
+        first_in = int(targets[0][1].shape[1])
+        self.proj_in = nn.Linear(input_dim, first_in)
+        prev_out = first_in
+        for name, w in targets:
+            w = w.contiguous()
+            # Zincir: onceki out != bu in ise ara projeksiyon
+            if prev_out != int(w.shape[1]):
+                layers.append(nn.Linear(prev_out, int(w.shape[1])))
+                layers.append(nn.GELU())
+            layers.append(FrozenLinearLoRA(w, rank=rank, name=name))
+            layers.append(nn.GELU())
+            prev_out = int(w.shape[0])
+        self.layers = layers
+        self.proj_out = nn.Linear(prev_out, out_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.proj_in(x)
+        for layer in self.layers:
+            h = layer(h)
+        return self.proj_out(h)
+
+    def lora_parameters(self) -> list[nn.Parameter]:
+        params: list[nn.Parameter] = []
+        for m in self.modules():
+            if isinstance(m, FrozenLinearLoRA):
+                params.extend([m.lora_A, m.lora_B])
+        return params