Upload scripts/job_quantize.py with huggingface_hub

scripts/job_quantize.py

@@ -0,0 +1,450 @@
+#!/usr/bin/env python3
+"""
+Fish Speech S2 Pro — Quantization Experiment (HF Job version)
+Downloads model, applies quantization at all phases, generates voice clone samples.
+"""
+import os, sys, json, time, gc, traceback
+import torch
+import torch.nn as nn
+import numpy as np
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+BASE_MODEL = "fishaudio/s2-pro"
+OUT = "/app/output"
+
+print("=== Fish Speech S2 Pro Quantization Experiment ===")
+print(f"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}")
+print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'none'}")
+print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB" if torch.cuda.is_available() else "")
+
+# Install deps
+print("\n[1/8] Installing dependencies...")
+os.system("pip install -q einops loguru ormsgpack hydra-core omegaconf safetensors torchaudio soundfile")
+os.system("pip install -q datasets")
+
+# Clone fish-speech
+if not os.path.exists("/app/fish-speech"):
+    print("\n[2/8] Cloning fish-speech repo...")
+    os.system("cd /app && git clone --depth 1 https://github.com/fishaudio/fish-speech.git")
+else:
+    print("\n[2/8] fish-speech already cloned")
+sys.path.insert(0, "/app/fish-speech")
+
+# Download model
+print("\n[3/8] Downloading S2 Pro model...")
+os.system(f"huggingface-cli download {BASE_MODEL} --local-dir /app/checkpoints/s2-pro")
+
+BASE_MODEL = "/app/checkpoints/s2-pro"
+
+# ============ QUANTIZATION CLASSES ============
+
+class FP8Linear(nn.Module):
+    def __init__(self, in_f, out_f, bias=True):
+        super().__init__()
+        self.in_features, self.out_features = in_f, out_f
+        self.register_buffer("weight", torch.empty(out_f, in_f, dtype=torch.float8_e4m3fn))
+        self.register_buffer("weight_scale", torch.empty(out_f, 1, dtype=torch.float32))
+        self.has_bias = bias
+        if bias: self.register_buffer("bias", torch.zeros(out_f, dtype=torch.bfloat16))
+        else: self.bias = None
+
+    @staticmethod
+    def from_linear(linear):
+        fp8 = FP8Linear(linear.in_features, linear.out_features, linear.bias is not None)
+        w = linear.weight.data.detach().bfloat16()
+        scale = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) / 448.0
+        w_q = (w / scale).round().clamp(-448, 448).to(torch.float8_e4m3fn)
+        fp8.weight.data.copy_(w_q)
+        fp8.weight_scale.data.copy_(scale)
+        if linear.bias is not None:
+            fp8.bias.data.copy_(linear.bias.data.detach().bfloat16())
+        return fp8
+
+    def forward(self, x):
+        return nn.functional.linear(x, self.weight.to(torch.bfloat16) * self.weight_scale, self.bias)
+
+class INT8Linear(nn.Module):
+    def __init__(self, in_f, out_f, bias=True):
+        super().__init__()
+        self.in_features, self.out_features = in_f, out_f
+        self.register_buffer("weight", torch.empty(out_f, in_f, dtype=torch.int8))
+        self.register_buffer("weight_scale", torch.empty(out_f, 1, dtype=torch.float32))
+        self.has_bias = bias
+        if bias: self.register_buffer("bias", torch.zeros(out_f, dtype=torch.bfloat16))
+        else: self.bias = None
+
+    @staticmethod
+    def from_linear(linear):
+        q = INT8Linear(linear.in_features, linear.out_features, linear.bias is not None)
+        w = linear.weight.data.detach().bfloat16()
+        scale = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) / 127.0
+        q.weight.data.copy_((w / scale).round().clamp(-128, 127).to(torch.int8))
+        q.weight_scale.data.copy_(scale)
+        if linear.bias is not None: q.bias.data.copy_(linear.bias.data.detach().bfloat16())
+        return q
+
+    def forward(self, x):
+        return nn.functional.linear(x, self.weight.to(torch.bfloat16) * self.weight_scale, self.bias)
+
+class INT4Linear(nn.Module):
+    def __init__(self, in_f, out_f, group_size=128, bias=True):
+        super().__init__()
+        self.in_features, self.out_features, self.group_size = in_f, out_f, group_size
+        self.register_buffer("weight_q", torch.empty(out_f, in_f, dtype=torch.int8))
+        self.register_buffer("weight_scale", torch.empty(out_f, (in_f + group_size - 1) // group_size, dtype=torch.float32))
+        self.has_bias = bias
+        if bias: self.register_buffer("bias", torch.zeros(out_f, dtype=torch.bfloat16))
+        else: self.bias = None
+
+    @staticmethod
+    def from_linear(linear, group_size=128):
+        in_f, out_f = linear.in_features, linear.out_features
+        q = INT4Linear(in_f, out_f, group_size, linear.bias is not None)
+        w = linear.weight.data.detach().bfloat16()
+        n_groups = (in_f + group_size - 1) // group_size
+        pad = n_groups * group_size - in_f
+        if pad > 0: w = nn.functional.pad(w, (0, pad))
+        w_g = w.reshape(out_f, n_groups, group_size)
+        scale = w_g.abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) / 7.0
+        q.weight_q.data.copy_((w_g / scale).round().clamp(-7, 7).to(torch.int8).reshape(out_f, -1)[:, :in_f])
+        q.weight_scale.data.copy_(scale.squeeze(-1)[:, :n_groups])
+        if linear.bias is not None: q.bias.data.copy_(linear.bias.data.detach().bfloat16())
+        return q
+
+    def forward(self, x):
+        s = self.weight_scale.repeat_interleave(self.group_size, dim=1)[:, :self.in_features]
+        return nn.functional.linear(x, self.weight_q[:, :self.in_features].to(torch.bfloat16) * s, self.bias)
+
+class INT3Linear(nn.Module):
+    def __init__(self, in_f, out_f, group_size=128, bias=True):
+        super().__init__()
+        self.in_features, self.out_features, self.group_size = in_f, out_f, group_size
+        self.register_buffer("weight_q", torch.empty(out_f, in_f, dtype=torch.int8))
+        self.register_buffer("weight_scale", torch.empty(out_f, (in_f + group_size - 1) // group_size, dtype=torch.float32))
+        self.has_bias = bias
+        if bias: self.register_buffer("bias", torch.zeros(out_f, dtype=torch.bfloat16))
+        else: self.bias = None
+
+    @staticmethod
+    def from_linear(linear, group_size=128):
+        in_f, out_f = linear.in_features, linear.out_features
+        q = INT3Linear(in_f, out_f, group_size, linear.bias is not None)
+        w = linear.weight.data.detach().bfloat16()
+        n_groups = (in_f + group_size - 1) // group_size
+        pad = n_groups * group_size - in_f
+        if pad > 0: w = nn.functional.pad(w, (0, pad))
+        w_g = w.reshape(out_f, n_groups, group_size)
+        scale = w_g.abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) / 3.0
+        q.weight_q.data.copy_((w_g / scale).round().clamp(-3, 3).to(torch.int8).reshape(out_f, -1)[:, :in_f])
+        q.weight_scale.data.copy_(scale.squeeze(-1)[:, :n_groups])
+        if linear.bias is not None: q.bias.data.copy_(linear.bias.data.detach().bfloat16())
+        return q
+
+    def forward(self, x):
+        s = self.weight_scale.repeat_interleave(self.group_size, dim=1)[:, :self.in_features]
+        return nn.functional.linear(x, self.weight_q[:, :self.in_features].to(torch.bfloat16) * s, self.bias)
+
+class INT2Linear(nn.Module):
+    def __init__(self, in_f, out_f, group_size=64, bias=True):
+        super().__init__()
+        self.in_features, self.out_features, self.group_size = in_f, out_f, group_size
+        self.register_buffer("weight_q", torch.empty(out_f, in_f, dtype=torch.int8))
+        self.register_buffer("weight_scale", torch.empty(out_f, (in_f + group_size - 1) // group_size, dtype=torch.float32))
+        self.has_bias = bias
+        if bias: self.register_buffer("bias", torch.zeros(out_f, dtype=torch.bfloat16))
+        else: self.bias = None
+
+    @staticmethod
+    def from_linear(linear, group_size=64):
+        in_f, out_f = linear.in_features, linear.out_features
+        q = INT2Linear(in_f, out_f, group_size, linear.bias is not None)
+        w = linear.weight.data.detach().bfloat16()
+        n_groups = (in_f + group_size - 1) // group_size
+        pad = n_groups * group_size - in_f
+        if pad > 0: w = nn.functional.pad(w, (0, pad))
+        w_g = w.reshape(out_f, n_groups, group_size)
+        scale = w_g.abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) / 1.0
+        q.weight_q.data.copy_((w_g / scale).round().clamp(-1, 1).to(torch.int8).reshape(out_f, -1)[:, :in_f])
+        q.weight_scale.data.copy_(scale.squeeze(-1)[:, :n_groups])
+        if linear.bias is not None: q.bias.data.copy_(linear.bias.data.detach().bfloat16())
+        return q
+
+    def forward(self, x):
+        s = self.weight_scale.repeat_interleave(self.group_size, dim=1)[:, :self.in_features]
+        return nn.functional.linear(x, self.weight_q[:, :self.in_features].to(torch.bfloat16) * s, self.bias)
+
+
+# ============ HELPERS ============
+
+def apply_quant(model, qcls, target="slow_ar", **kw):
+    count = 0
+    skip = ['embed', 'norm']
+    for name, mod in list(model.named_modules()):
+        if not isinstance(mod, nn.Linear): continue
+        if any(s in name for s in skip): continue
+        if target == "slow_ar" and "fast_" in name: continue
+        parts = name.split(".")
+        parent = model
+        for p in parts[:-1]: parent = getattr(parent, p)
+        try:
+            setattr(parent, parts[-1], qcls.from_linear(mod, **kw))
+            count += 1
+        except: pass
+    return model, count
+
+def model_size_mb(m):
+    t = sum(p.numel() * p.element_size() for p in m.parameters())
+    t += sum(b.numel() * b.element_size() for b in m.buffers())
+    return t / (1024*1024)
+
+def generate_sample(model, codec, text, out_path):
+    """Generate TTS sample"""
+    import soundfile as sf
+    from fish_speech.models.text2semantic.inference import generate, decode_one_token_ar
+    from fish_speech.content_sequence import TextPart
+    from fish_speech.conversation import Conversation, Message
+
+    try:
+        conv = Conversation()
+        conv.add_message(Message(role="user", parts=[TextPart(text="")]))
+        conv.add_message(Message(role="assistant", parts=[TextPart(text=text)]))
+        prompt = conv.encode_for_inference(model.config)
+        cd = 1 + model.config.num_codebooks
+        am = torch.zeros(1, cd, prompt.shape[-1], dtype=torch.bool, device=DEVICE)
+        ap = torch.zeros(1, cd, prompt.shape[-1], dtype=torch.long, device=DEVICE)
+
+        if not getattr(model, '_cache_done', False):
+            model.setup_caches(1, model.config.max_seq_len, dtype=DTYPE)
+            model._cache_done = True
+
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            result = generate(model=model, prompt=prompt, max_new_tokens=512,
+                audio_masks=am, audio_parts=ap, temperature=0.7, top_p=0.7, top_k=30,
+                decode_one_token=decode_one_token_ar)
+
+        codes = result[0:1, :, :].unsqueeze(0)
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            audio = codec.decode(codes.to(DEVICE))
+        audio_np = audio.squeeze().cpu().float().numpy()
+        sr = getattr(codec, 'sample_rate', 44100)
+        sf.write(out_path, audio_np, sr)
+        dur = len(audio_np) / sr
+        print(f"  Saved {out_path} ({dur:.1f}s)")
+        return True, dur
+    except Exception as e:
+        print(f"  Sample gen failed: {e}")
+        traceback.print_exc()
+        return False, 0
+
+def generate_clone(model, codec, text, ref_path, ref_text, out_path):
+    """Voice clone sample"""
+    import torchaudio, soundfile as sf
+    from fish_speech.models.text2semantic.inference import generate, decode_one_token_ar
+    from fish_speech.content_sequence import TextPart, VQPart
+    from fish_speech.conversation import Conversation, Message
+
+    try:
+        wav, sr = torchaudio.load(ref_path)
+        if wav.shape[0] > 1: wav = wav.mean(dim=0, keepdim=True)
+        if sr != 44100: wav = torchaudio.functional.resample(wav, sr, 44100)
+        wav = wav.to(DEVICE)
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            enc = codec.encode(wav.unsqueeze(0))
+            ptokens = (enc[0] if isinstance(enc, tuple) else enc).cpu().numpy()
+
+        conv = Conversation()
+        conv.add_message(Message(role="user", parts=[VQPart(codes=ptokens), TextPart(text=ref_text)]))
+        conv.add_message(Message(role="assistant", parts=[TextPart(text=text)]))
+        prompt = conv.encode_for_inference(model.config)
+        cd = 1 + model.config.num_codebooks
+        am = torch.zeros(1, cd, prompt.shape[-1], dtype=torch.bool, device=DEVICE)
+        ap = torch.zeros(1, cd, prompt.shape[-1], dtype=torch.long, device=DEVICE)
+
+        if not getattr(model, '_cache_done', False):
+            model.setup_caches(1, model.config.max_seq_len, dtype=DTYPE)
+            model._cache_done = True
+
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            result = generate(model=model, prompt=prompt, max_new_tokens=512,
+                audio_masks=am, audio_parts=ap, temperature=0.7, top_p=0.7, top_k=30,
+                decode_one_token=decode_one_token_ar)
+        codes = result[0:1, :, :].unsqueeze(0)
+        with torch.autocast(device_type="cuda", dtype=DTYPE):
+            audio = codec.decode(codes.to(DEVICE))
+        audio_np = audio.squeeze().cpu().float().numpy()
+        sr = getattr(codec, 'sample_rate', 44100)
+        sf.write(out_path, audio_np, sr)
+        dur = len(audio_np) / sr
+        print(f"  Clone saved {out_path} ({dur:.1f}s)")
+        return True, dur
+    except Exception as e:
+        print(f"  Clone failed: {e}")
+        traceback.print_exc()
+        return False, 0
+
+
+# ============ PHASE RUNNER ============
+
+def run_phase(pid, qcls, target, codec, ref_path, ref_text, test_text, clone_text, **kw):
+    from fish_speech.models.text2semantic.inference import init_model
+    from safetensors.torch import save_file
+
+    phase_dir = f"{OUT}/{pid}"
+    os.makedirs(phase_dir, exist_ok=True)
+    os.makedirs(f"{OUT}/samples", exist_ok=True)
+
+    print(f"\n{'='*60}")
+    print(f"  {pid.upper()}: {qcls.__name__} ({target})")
+    print(f"{'='*60}")
+
+    model, _ = init_model(BASE_MODEL, DEVICE, DTYPE, compile=False)
+    orig = model_size_mb(model)
+
+    t0 = time.time()
+    model, nl = apply_quant(model, qcls, target=target, **kw)
+    model = model.to(DEVICE)
+    t_q = time.time() - t0
+    qs = model_size_mb(model)
+    ratio = orig / qs if qs > 0 else 0
+    print(f"  {orig:.0f} -> {qs:.0f} MB ({ratio:.2f}x, {nl} layers, {t_q:.1f}s)")
+
+    sp = f"{phase_dir}/model.safetensors"
+    save_file(model.state_dict(), sp)
+    disk = os.path.getsize(sp) / (1024*1024)
+    print(f"  Disk: {disk:.0f} MB")
+
+    tts_ok, tts_d = generate_sample(model, codec, test_text, f"{OUT}/samples/{pid}_tts.wav")
+    clone_ok, clone_d = False, 0
+    if ref_path and os.path.exists(ref_path):
+        clone_ok, clone_d = generate_clone(model, codec, clone_text, ref_path, ref_text, f"{OUT}/samples/{pid}_clone.wav")
+
+    del model; gc.collect(); torch.cuda.empty_cache()
+
+    r = {"phase": pid, "method": qcls.__name__, "target": target,
+         "orig_mb": round(orig), "quant_mb": round(qs), "disk_mb": round(disk),
+         "ratio": round(ratio, 3), "layers": nl, "time_s": round(t_q,1),
+         "tts_ok": tts_ok, "tts_d": round(tts_d,1),
+         "clone_ok": clone_ok, "clone_d": round(clone_d,1)}
+    with open(f"{phase_dir}/results.json","w") as f: json.dump(r,f,indent=2)
+    return r
+
+# ============ MAIN ============
+
+TEST_TEXT = "The quick brown fox jumps over the lazy dog. Artificial intelligence is transforming how we communicate."
+CLONE_TEXT = "Hello everyone, welcome to this special presentation. Today we explore the fascinating world of neural text to speech and voice cloning technology."
+REF_TEXT = "This is a reference voice recording."
+CELEB_TEXT = "Good morning. I want to tell you something about the universe. Every atom in your body came from a star that exploded. We are all made of star stuff."
+
+PHASES = {
+    "1a": (FP8Linear, "slow_ar", {}),
+    "1b": (INT4Linear, "slow_ar", {"group_size": 128}),
+    "2a": (INT4Linear, "all", {"group_size": 128}),
+    "2b": (INT8Linear, "slow_ar", {}),
+    "2c": (INT3Linear, "slow_ar", {"group_size": 128}),
+    "3a": (INT2Linear, "slow_ar", {"group_size": 64}),
+    "3b": (INT2Linear, "all", {"group_size": 64}),
+}
+
+def main():
+    os.makedirs(f"{OUT}/samples", exist_ok=True)
+    from fish_speech.models.text2semantic.inference import init_model
+    from fish_speech.models.dac.inference import load_codec_model
+
+    print("\n[4/8] Loading base model...")
+    model_base, _ = init_model(BASE_MODEL, DEVICE, DTYPE, compile=False)
+    orig = model_size_mb(model_base)
+    print(f"  Base model: {orig:.0f} MB ({sum(p.numel() for p in model_base.parameters())/1e9:.2f}B params)")
+
+    print("\n[5/8] Loading codec...")
+    codec = load_codec_model(f"{BASE_MODEL}/codec.pth", DEVICE, DTYPE)
+
+    # Generate reference + baseline
+    ref_path = f"{OUT}/reference_celebrity.wav"
+    print("\n[6/8] Generating reference audio & baseline...")
+    try:
+        generate_sample(model_base, codec, CELEB_TEXT, ref_path)
+    except Exception as e:
+        print(f"  Ref gen warning: {e}")
+        ref_path = None
+
+    try:
+        generate_sample(model_base, codec, TEST_TEXT, f"{OUT}/samples/baseline_bf16_tts.wav")
+    except: pass
+
+    if ref_path:
+        try:
+            generate_clone(model_base, codec, CLONE_TEXT, ref_path, REF_TEXT, f"{OUT}/samples/baseline_bf16_clone.wav")
+        except: pass
+
+    del model_base; gc.collect(); torch.cuda.empty_cache()
+
+    # Run all phases
+    print("\n[7/8] Running quantization phases...")
+    all_r = [{"phase": "baseline_bf16", "orig_mb": round(orig), "quant_mb": round(orig), "disk_mb": round(orig), "ratio": 1.0}]
+
+    for pid, (qcls, target, kw) in PHASES.items():
+        try:
+            r = run_phase(f"phase{pid}", qcls, target, codec, ref_path, REF_TEXT, TEST_TEXT, CLONE_TEXT, **kw)
+            all_r.append(r)
+        except Exception as e:
+            print(f"Phase {pid} FAILED: {e}")
+            traceback.print_exc()
+            all_r.append({"phase": f"phase{pid}", "error": str(e)})
+
+    # Summary
+    print(f"\n{'='*70}")
+    print("  RESULTS SUMMARY")
+    print(f"{'='*70}")
+    print(f"{'Phase':<14} {'Method':<14} {'Disk MB':<10} {'Ratio':<8} {'TTS':<5} {'Clone':<5}")
+    print("-"*60)
+    for r in all_r:
+        print(f"{r.get('phase','?'):<14} {r.get('method','bf16'):<14} {r.get('disk_mb','?'):<10} {r.get('ratio',1):<8.2f} {'OK' if r.get('tts_ok') else 'FAIL':<5} {'OK' if r.get('clone_ok') else 'FAIL':<5}")
+
+    with open(f"{OUT}/all_results.json","w") as f: json.dump(all_r,f,indent=2)
+
+    # Upload to Hub
+    print("\n[8/8] Uploading results to HuggingFace Hub...")
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        repo = "Swagcrew/fish-speech-s2-quantized"
+        api.create_repo(repo_id=repo, repo_type="model", exist_ok=True, private=False)
+
+        # Upload all results
+        api.upload_file(path_or_fileobj=f"{OUT}/all_results.json", path_in_repo="all_results.json", repo_id=repo, repo_type="model")
+
+        # Upload samples
+        samples_dir = f"{OUT}/samples"
+        if os.path.exists(samples_dir):
+            for fn in os.listdir(samples_dir):
+                if fn.endswith(".wav"):
+                    api.upload_file(path_or_fileobj=os.path.join(samples_dir, fn),
+                                    path_in_repo=f"samples/{fn}", repo_id=repo, repo_type="model")
+
+        # Upload individual phase results
+        for pid in PHASES:
+            phase_dir = f"{OUT}/phase{pid}"
+            if os.path.exists(f"{phase_dir}/results.json"):
+                api.upload_file(f"{phase_dir}/results.json", f"phase{pid}/results.json", repo, repo_type="model")
+            if os.path.exists(f"{phase_dir}/model.safetensors"):
+                api.upload_file(f"{phase_dir}/model.safetensors", f"phase{pid}/model.safetensors", repo, repo_type="model")
+
+        # Upload README
+        readme = """# Fish Speech S2 Pro — Quantization Experiments
+Multi-phase quantization with voice cloning samples.
+See all_results.json for details and samples/ for audio.
+"""
+        api.upload_file(path_or_fileobj=readme.encode(), path_in_repo="README.md", repo_id=repo, repo_type="model")
+
+        print(f"  Uploaded to https://huggingface.co/{repo}")
+    except Exception as e:
+        print(f"  Upload failed: {e}")
+        traceback.print_exc()
+
+    print("\nDONE! All phases complete.")
+
+if __name__ == "__main__":
+    main()
+