533 lines
19 KiB
Python
533 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
SDXL Long Context Converter — Learned Sinusoidal Extrapolation (Optimized)
|
|
|
|
Extends positional embeddings by fitting parametric sinusoidal models
|
|
to each dimension, matching the mathematical structure learned during
|
|
CLIP training.
|
|
|
|
Methods:
|
|
learned_sinusoidal — fit N sinusoids + trend
|
|
vaswani_fitted — fit Vaswani PE formula
|
|
hybrid — best of all per dimension (recommended)
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
try:
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
except ImportError:
|
|
print("ERROR: pip install safetensors")
|
|
sys.exit(1)
|
|
|
|
try:
|
|
from scipy.optimize import curve_fit
|
|
except ImportError:
|
|
print("Installing scipy...")
|
|
import subprocess
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", "scipy"])
|
|
from scipy.optimize import curve_fit
|
|
|
|
|
|
POS_EMB_KEYS = [
|
|
"conditioner.embedders.0.transformer.text_model.embeddings.position_embedding.weight",
|
|
"cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
|
|
"text_encoder.text_model.embeddings.position_embedding.weight",
|
|
"conditioner.embedders.1.model.positional_embedding",
|
|
"conditioner.embedders.1.transformer.text_model.embeddings.position_embedding.weight",
|
|
"text_encoder_2.text_model.embeddings.position_embedding.weight",
|
|
]
|
|
|
|
MAX_LENGTH = 248
|
|
|
|
|
|
# ================================================================
|
|
# Progress bar
|
|
# ================================================================
|
|
|
|
def print_progress(current, total, prefix="", width=40):
|
|
pct = current / total
|
|
filled = int(width * pct)
|
|
bar = "█" * filled + "░" * (width - filled)
|
|
print(f"\r {prefix} [{bar}] {current}/{total} ({pct*100:.0f}%)", end="", flush=True)
|
|
|
|
|
|
# ================================================================
|
|
# Vectorized spectral fallback (fast)
|
|
# ================================================================
|
|
|
|
def spectral_extrapolate_single(signal, target_length, damping=0.08):
|
|
"""Fast spectral extrapolation for a single dimension."""
|
|
orig_len = len(signal)
|
|
dc = np.mean(signal)
|
|
spectrum = np.fft.rfft(signal - dc)
|
|
n_freqs = len(spectrum)
|
|
|
|
t_new = np.arange(target_length, dtype=np.float64) / orig_len
|
|
|
|
# Vectorized: compute all frequencies at once
|
|
amps = np.abs(spectrum)
|
|
phases = np.angle(spectrum)
|
|
|
|
# Normalization factors
|
|
norms = np.full(n_freqs, 2.0 / orig_len)
|
|
norms[0] = 1.0 / orig_len
|
|
if orig_len % 2 == 0:
|
|
norms[-1] = 1.0 / orig_len
|
|
|
|
# Damping envelope for extrapolated region
|
|
extrap_dist = np.zeros(target_length, dtype=np.float64)
|
|
extrap_dist[orig_len:] = np.arange(1, target_length - orig_len + 1, dtype=np.float64) / orig_len
|
|
|
|
reconstructed = np.zeros(target_length, dtype=np.float64)
|
|
|
|
# Vectorize over positions (not frequencies) for better cache behavior
|
|
k_values = np.arange(n_freqs, dtype=np.float64)
|
|
freq_ratios = k_values / n_freqs
|
|
|
|
for k in range(n_freqs):
|
|
if amps[k] < 1e-10:
|
|
continue
|
|
damp = np.exp(-damping * freq_ratios[k] * freq_ratios[k] * extrap_dist * orig_len)
|
|
reconstructed += amps[k] * norms[k] * np.cos(2 * np.pi * k * t_new - phases[k]) * damp
|
|
|
|
return reconstructed + dc
|
|
|
|
|
|
def spectral_extrapolate_batch(signals, target_length, damping=0.08):
|
|
"""Batch spectral extrapolation for multiple dimensions at once."""
|
|
n_dims, orig_len = signals.shape
|
|
result = np.zeros((n_dims, target_length), dtype=np.float64)
|
|
|
|
for d in range(n_dims):
|
|
result[d] = spectral_extrapolate_single(signals[d], target_length, damping)
|
|
|
|
return result
|
|
|
|
|
|
# ================================================================
|
|
# Sinusoidal model fitting (optimized)
|
|
# ================================================================
|
|
|
|
def sinusoidal_model_2(pos, a, b, A1, f1, p1, A2, f2, p2):
|
|
return (a + b * pos +
|
|
A1 * np.sin(2 * np.pi * f1 * pos + p1) +
|
|
A2 * np.sin(2 * np.pi * f2 * pos + p2))
|
|
|
|
|
|
def sinusoidal_model_3(pos, a, b, A1, f1, p1, A2, f2, p2, A3, f3, p3):
|
|
return (a + b * pos +
|
|
A1 * np.sin(2 * np.pi * f1 * pos + p1) +
|
|
A2 * np.sin(2 * np.pi * f2 * pos + p2) +
|
|
A3 * np.sin(2 * np.pi * f3 * pos + p3))
|
|
|
|
|
|
def fit_sinusoidal_fast(signal, positions, n_components=3):
|
|
"""
|
|
Fast sinusoidal fitting with reduced maxfev and smart initialization.
|
|
Returns (params, model_func, r_squared).
|
|
"""
|
|
orig_len = len(signal)
|
|
dc = np.mean(signal)
|
|
trend = (signal[-1] - signal[0]) / orig_len if orig_len > 1 else 0.0
|
|
|
|
spectrum = np.fft.rfft(signal - dc)
|
|
amps = np.abs(spectrum)
|
|
top_freqs = np.argsort(amps[1:])[::-1][:n_components] + 1
|
|
|
|
if n_components <= 2:
|
|
model_func = sinusoidal_model_2
|
|
n_comp = 2
|
|
else:
|
|
model_func = sinusoidal_model_3
|
|
n_comp = 3
|
|
|
|
p0 = [dc, trend]
|
|
for i in range(n_comp):
|
|
idx = top_freqs[i] if i < len(top_freqs) else 1
|
|
p0.extend([
|
|
amps[idx] * 2 / orig_len,
|
|
idx / orig_len,
|
|
0.0
|
|
])
|
|
|
|
try:
|
|
popt, _ = curve_fit(model_func, positions, signal, p0=p0, maxfev=2000)
|
|
fitted = model_func(positions, *popt)
|
|
residual = np.mean((signal - fitted) ** 2)
|
|
variance = np.var(signal)
|
|
r2 = 1.0 - residual / (variance + 1e-10) if variance > 1e-10 else 1.0
|
|
return popt, model_func, r2
|
|
except (RuntimeError, ValueError):
|
|
return None, model_func, 0.0
|
|
|
|
|
|
def fit_vaswani_fast(signal, dim_idx, d_model, positions):
|
|
"""
|
|
Fast Vaswani formula fitting with reduced iterations.
|
|
Returns (params, model_func, r_squared).
|
|
"""
|
|
dc = np.mean(signal)
|
|
signal_ac = signal - dc
|
|
amp_init = np.std(signal_ac) * np.sqrt(2)
|
|
|
|
def model(pos, amplitude, base, phase, offset):
|
|
freq = 1.0 / np.power(np.abs(base) + 1.0, 2.0 * dim_idx / d_model)
|
|
return offset + amplitude * np.sin(pos * freq + phase)
|
|
|
|
try:
|
|
popt, _ = curve_fit(
|
|
model, positions, signal,
|
|
p0=[amp_init, 10000.0, 0.0, dc],
|
|
maxfev=2000,
|
|
bounds=(
|
|
[-np.inf, 1.0, -2 * np.pi, -np.inf],
|
|
[np.inf, 100000.0, 2 * np.pi, np.inf]
|
|
)
|
|
)
|
|
fitted = model(positions, *popt)
|
|
residual = np.mean((signal - fitted) ** 2)
|
|
variance = np.var(signal)
|
|
r2 = 1.0 - residual / (variance + 1e-10) if variance > 1e-10 else 1.0
|
|
return popt, model, r2
|
|
except (RuntimeError, ValueError):
|
|
return None, model, 0.0
|
|
|
|
|
|
# ================================================================
|
|
# Extension methods
|
|
# ================================================================
|
|
|
|
def quintic_step(t):
|
|
t = np.clip(t, 0.0, 1.0)
|
|
return t * t * t * (t * (t * 6 - 15) + 10)
|
|
|
|
|
|
def preserve_and_blend(result, original, blend_width=12):
|
|
orig_len = original.shape[0]
|
|
if blend_width <= 0:
|
|
result[:orig_len] = original
|
|
return result
|
|
|
|
blend_start = max(0, orig_len - blend_width)
|
|
result[:blend_start] = original[:blend_start]
|
|
|
|
for i in range(blend_start, orig_len):
|
|
alpha = quintic_step((i - blend_start) / blend_width)
|
|
result[i] = (1.0 - alpha) * original[i] + alpha * result[i]
|
|
|
|
return result
|
|
|
|
|
|
def match_norms(result, original, start_idx):
|
|
orig_norms = np.linalg.norm(original, axis=1)
|
|
mean_norm = np.mean(orig_norms)
|
|
|
|
for i in range(start_idx, result.shape[0]):
|
|
norm = np.linalg.norm(result[i])
|
|
if norm > 1e-10:
|
|
result[i] *= mean_norm / norm
|
|
|
|
return result
|
|
|
|
|
|
def extend_learned_sinusoidal(pos_emb, target_length, n_components=3,
|
|
blend_width=12):
|
|
"""Fit N sinusoids + trend to each dimension."""
|
|
orig_len, dim = pos_emb.shape
|
|
pos_float = pos_emb.float().numpy().astype(np.float64)
|
|
result = np.zeros((target_length, dim), dtype=np.float64)
|
|
|
|
positions_orig = np.arange(orig_len, dtype=np.float64) / orig_len
|
|
positions_ext = np.arange(target_length, dtype=np.float64) / orig_len
|
|
|
|
fit_qualities = []
|
|
t_start = time.time()
|
|
|
|
for d in range(dim):
|
|
if d % 50 == 0 or d == dim - 1:
|
|
print_progress(d + 1, dim, "Fitting sinusoids")
|
|
|
|
signal = pos_float[:, d]
|
|
params, model_func, r2 = fit_sinusoidal_fast(signal, positions_orig, n_components)
|
|
fit_qualities.append(r2)
|
|
|
|
if params is not None and r2 > 0.3:
|
|
result[:, d] = model_func(positions_ext, *params)
|
|
else:
|
|
result[:, d] = spectral_extrapolate_single(signal, target_length)
|
|
|
|
elapsed = time.time() - t_start
|
|
print(f"\r Fitting sinusoids: done in {elapsed:.1f}s" + " " * 30)
|
|
|
|
result = preserve_and_blend(result, pos_float, blend_width)
|
|
result = match_norms(result, pos_float, orig_len)
|
|
|
|
fq = np.array(fit_qualities)
|
|
print(f" Fit quality (R²): mean={np.mean(fq):.4f}, "
|
|
f"min={np.min(fq):.4f}, max={np.max(fq):.4f}")
|
|
print(f" Good fits (R²>0.9): {np.sum(fq > 0.9)}/{dim}")
|
|
print(f" Fair fits (R²>0.5): {np.sum(fq > 0.5)}/{dim}")
|
|
print(f" Poor fits (R²<0.3): {np.sum(fq < 0.3)}/{dim}")
|
|
|
|
return torch.from_numpy(result).to(pos_emb.dtype)
|
|
|
|
|
|
def extend_vaswani_fitted(pos_emb, target_length, blend_width=12):
|
|
"""Fit Vaswani PE formula to each dimension."""
|
|
orig_len, dim = pos_emb.shape
|
|
pos_float = pos_emb.float().numpy().astype(np.float64)
|
|
result = np.zeros((target_length, dim), dtype=np.float64)
|
|
|
|
positions_orig = np.arange(orig_len, dtype=np.float64)
|
|
positions_ext = np.arange(target_length, dtype=np.float64)
|
|
|
|
good_fits = 0
|
|
fallback_count = 0
|
|
t_start = time.time()
|
|
|
|
for d in range(dim):
|
|
if d % 50 == 0 or d == dim - 1:
|
|
print_progress(d + 1, dim, "Fitting Vaswani")
|
|
|
|
signal = pos_float[:, d]
|
|
params, model, r2 = fit_vaswani_fast(signal, d, dim, positions_orig)
|
|
|
|
if params is not None and r2 > 0.5:
|
|
result[:, d] = model(positions_ext, *params)
|
|
good_fits += 1
|
|
else:
|
|
result[:, d] = spectral_extrapolate_single(signal, target_length)
|
|
fallback_count += 1
|
|
|
|
elapsed = time.time() - t_start
|
|
print(f"\r Fitting Vaswani: done in {elapsed:.1f}s" + " " * 30)
|
|
print(f" Vaswani fits: {good_fits}/{dim} good, {fallback_count}/{dim} spectral fallback")
|
|
|
|
result = preserve_and_blend(result, pos_float, blend_width)
|
|
result = match_norms(result, pos_float, orig_len)
|
|
return torch.from_numpy(result).to(pos_emb.dtype)
|
|
|
|
|
|
def extend_hybrid(pos_emb, target_length, blend_width=12):
|
|
"""
|
|
Try Vaswani first, then sinusoidal, then spectral fallback.
|
|
Pick best R² for each dimension.
|
|
"""
|
|
orig_len, dim = pos_emb.shape
|
|
pos_float = pos_emb.float().numpy().astype(np.float64)
|
|
result = np.zeros((target_length, dim), dtype=np.float64)
|
|
|
|
positions_orig_norm = np.arange(orig_len, dtype=np.float64) / orig_len
|
|
positions_ext_norm = np.arange(target_length, dtype=np.float64) / orig_len
|
|
positions_orig_abs = np.arange(orig_len, dtype=np.float64)
|
|
positions_ext_abs = np.arange(target_length, dtype=np.float64)
|
|
|
|
method_counts = {"vaswani": 0, "sinusoidal": 0, "spectral": 0}
|
|
r2_threshold = 0.7
|
|
t_start = time.time()
|
|
|
|
for d in range(dim):
|
|
if d % 50 == 0 or d == dim - 1:
|
|
elapsed = time.time() - t_start
|
|
eta = (elapsed / (d + 1)) * (dim - d - 1) if d > 0 else 0
|
|
print_progress(d + 1, dim, f"Hybrid (ETA {eta:.0f}s)")
|
|
|
|
signal = pos_float[:, d]
|
|
|
|
best_r2 = -1.0
|
|
best_result = None
|
|
best_method = "spectral"
|
|
|
|
# Try Vaswani
|
|
params_v, model_v, r2_v = fit_vaswani_fast(signal, d, dim, positions_orig_abs)
|
|
if params_v is not None and r2_v > best_r2:
|
|
best_r2 = r2_v
|
|
best_result = model_v(positions_ext_abs, *params_v)
|
|
best_method = "vaswani"
|
|
|
|
# Try sinusoidal (only if Vaswani didn't reach threshold)
|
|
if best_r2 < r2_threshold:
|
|
params_s, model_s, r2_s = fit_sinusoidal_fast(
|
|
signal, positions_orig_norm, n_components=3
|
|
)
|
|
if params_s is not None and r2_s > best_r2:
|
|
best_r2 = r2_s
|
|
best_result = model_s(positions_ext_norm, *params_s)
|
|
best_method = "sinusoidal"
|
|
|
|
# Spectral fallback if both poor
|
|
if best_r2 < r2_threshold:
|
|
best_result = spectral_extrapolate_single(signal, target_length)
|
|
best_method = "spectral"
|
|
|
|
result[:, d] = best_result
|
|
method_counts[best_method] += 1
|
|
|
|
elapsed = time.time() - t_start
|
|
print(f"\r Hybrid fitting: done in {elapsed:.1f}s" + " " * 40)
|
|
print(f" Method selection: {method_counts}")
|
|
total = sum(method_counts.values())
|
|
for m, c in method_counts.items():
|
|
print(f" {m}: {c}/{total} ({c/total*100:.1f}%)")
|
|
|
|
result = preserve_and_blend(result, pos_float, blend_width)
|
|
result = match_norms(result, pos_float, orig_len)
|
|
return torch.from_numpy(result).to(pos_emb.dtype)
|
|
|
|
|
|
# ================================================================
|
|
# Analysis
|
|
# ================================================================
|
|
|
|
def analyze_embedding(name, original, extended, method_name):
|
|
orig_len = original.shape[0]
|
|
orig_np = original.float().numpy()
|
|
ext_np = extended.float().numpy()
|
|
|
|
diff = np.abs(orig_np - ext_np[:orig_len])
|
|
|
|
print(f"\n [{method_name}] {name}")
|
|
print(f" Shape: {tuple(original.shape)} -> {tuple(extended.shape)}")
|
|
print(f" Original preserved: max_diff={np.max(diff):.2e}, mean_diff={np.mean(diff):.2e}")
|
|
|
|
orig_norms = np.linalg.norm(orig_np, axis=1)
|
|
ext_norms = np.linalg.norm(ext_np[orig_len:], axis=1)
|
|
|
|
print(f" Original norms: mean={np.mean(orig_norms):.4f} "
|
|
f"std={np.std(orig_norms):.4f} "
|
|
f"[{np.min(orig_norms):.4f}, {np.max(orig_norms):.4f}]")
|
|
print(f" Extended norms: mean={np.mean(ext_norms):.4f} "
|
|
f"std={np.std(ext_norms):.4f} "
|
|
f"[{np.min(ext_norms):.4f}, {np.max(ext_norms):.4f}]")
|
|
|
|
ratio = np.mean(ext_norms) / np.mean(orig_norms)
|
|
quality = ("EXCELLENT" if 0.95 <= ratio <= 1.05 else
|
|
"GOOD" if 0.9 <= ratio <= 1.1 else
|
|
"ACCEPTABLE" if 0.8 <= ratio <= 1.2 else "WARNING")
|
|
print(f" Norm ratio: {ratio:.4f} — {quality}")
|
|
|
|
last_orig = orig_np[-1]
|
|
first_ext = ext_np[orig_len]
|
|
cos_sim = np.dot(last_orig, first_ext) / (
|
|
np.linalg.norm(last_orig) * np.linalg.norm(first_ext) + 1e-10)
|
|
print(f" Boundary cos_sim: {cos_sim:.6f}")
|
|
|
|
inter_dists = [np.linalg.norm(orig_np[i + 1] - orig_np[i])
|
|
for i in range(orig_len - 1)]
|
|
boundary_dist = np.linalg.norm(last_orig - first_ext)
|
|
print(f" Boundary/avg_step: {boundary_dist / (np.mean(inter_dists) + 1e-10):.4f}")
|
|
|
|
|
|
# ================================================================
|
|
# Main
|
|
# ================================================================
|
|
|
|
ALL_METHODS = ["learned_sinusoidal", "vaswani_fitted", "hybrid"]
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="SDXL Long Context — Learned Sinusoidal Extrapolation",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Methods:
|
|
learned_sinusoidal Fit N sinusoids + trend to each dimension
|
|
vaswani_fitted Fit original Transformer PE formula
|
|
hybrid Best of all methods per dimension (recommended)
|
|
"""
|
|
)
|
|
parser.add_argument("--input", required=True)
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--target_len", type=int, default=MAX_LENGTH)
|
|
parser.add_argument("--method", choices=ALL_METHODS, default="hybrid")
|
|
parser.add_argument("--n_components", type=int, default=3)
|
|
parser.add_argument("--blend_width", type=int, default=12)
|
|
parser.add_argument("--analyze", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
target = min(args.target_len, MAX_LENGTH)
|
|
|
|
print(f"SDXL Long Context — Learned Sinusoidal Extrapolation")
|
|
print(f"=" * 55)
|
|
print(f"Input: {args.input}")
|
|
print(f"Output: {args.output}")
|
|
print(f"Target: {target} tokens")
|
|
print(f"Method: {args.method}")
|
|
if args.method == "learned_sinusoidal":
|
|
print(f"Components: {args.n_components}")
|
|
print(f"Blend width: {args.blend_width}")
|
|
|
|
print(f"\nLoading...")
|
|
state_dict = {}
|
|
with safe_open(args.input, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
state_dict[key] = f.get_tensor(key)
|
|
print(f" {len(state_dict)} tensors")
|
|
|
|
modified = []
|
|
for key in POS_EMB_KEYS:
|
|
if key not in state_dict:
|
|
continue
|
|
original = state_dict[key]
|
|
if original.dim() != 2 or original.shape[0] >= target:
|
|
continue
|
|
|
|
print(f"\nExtending: {key}")
|
|
print(f" {tuple(original.shape)} -> ({target}, {original.shape[1]})")
|
|
|
|
t0 = time.time()
|
|
|
|
if args.method == "learned_sinusoidal":
|
|
extended = extend_learned_sinusoidal(
|
|
original, target, args.n_components, args.blend_width)
|
|
elif args.method == "vaswani_fitted":
|
|
extended = extend_vaswani_fitted(original, target, args.blend_width)
|
|
elif args.method == "hybrid":
|
|
extended = extend_hybrid(original, target, args.blend_width)
|
|
|
|
elapsed = time.time() - t0
|
|
print(f" Total time: {elapsed:.1f}s")
|
|
|
|
if args.analyze:
|
|
analyze_embedding(key, original, extended, args.method)
|
|
|
|
state_dict[key] = extended
|
|
modified.append(key)
|
|
|
|
if not modified:
|
|
print("\nERROR: No positional embeddings found!")
|
|
for key in state_dict:
|
|
if "position" in key.lower():
|
|
print(f" {key}: {tuple(state_dict[key].shape)}")
|
|
sys.exit(1)
|
|
|
|
print(f"\nSaving: {args.output}")
|
|
metadata = {
|
|
"long_context": f"learned_{args.method}",
|
|
"original_length": "77",
|
|
"target_length": str(target),
|
|
"method": args.method,
|
|
"modified_keys": json.dumps(modified),
|
|
}
|
|
clean = {k: v.contiguous() for k, v in state_dict.items()}
|
|
save_file(clean, args.output, metadata=metadata)
|
|
|
|
size = os.path.getsize(args.output) / (1024 ** 3)
|
|
print(f"\nDone! ({size:.2f} GB)")
|
|
print(f"Modified: {len(modified)} tensors")
|
|
for key in modified:
|
|
print(f" - {key}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|