File size: 5,545 Bytes
c6b84b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | #!/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
|