import os import gc import torch import shutil import uuid import gradio as gr from huggingface_hub import HfApi, hf_hub_download from safetensors.torch import load_file, save_file # Architecture profiles tuned specifically for HeartMuLa variations ARCH_PROFILES = { "HeartMuLa / Transformer Core (3B / RL / Gen)": [ "audio_embeddings", "audio_head", "sa_norm", "mlp_norm", "embed", "norm" ], "HeartMuLa / Flow Matching Codec": [ "cond_feature_emb", "timestep_embedder", "adaln_single", "connection_proj", "norm" ], "None / Standard (Quantize All Layers)": [] } def convert_and_upload(token, source_repo, target_repo, precision, target_components, arch_profile): if not token: yield "❌ Error: Please provide a valid Hugging Face Write Token." return if not target_repo.strip() or "/" not in target_repo: yield "❌ Error: Target Repository must be in format 'username/repo-name'." return if not target_components: yield "❌ Error: Please select at least one component to process." return # Map precision types target_dtype = None is_int8 = precision == "INT8" is_int4 = precision == "INT4" if precision == "FP16": target_dtype = torch.float16 elif precision == "BF16": target_dtype = torch.bfloat16 api = HfApi(token=token) yield f"🔄 Verifying target repo: {target_repo}..." try: api.create_repo(repo_id=target_repo, exist_ok=True, private=False) except Exception as e: yield f"❌ Error creating repo: {str(e)}" return yield f"📋 Fetching files from {source_repo}..." try: files = api.list_repo_files(source_repo) except Exception as e: yield f"❌ Error fetching files: {str(e)}" return cache_dir = f"./hf_cache_{uuid.uuid4().hex[:8]}" success_count, error_count = 0, 0 exclude_prefixes = ARCH_PROFILES.get(arch_profile, []) for file in files: # Check if file matches folder-level filtering or root-level targeting in_target_folder = any(f"{comp}/" in file for comp in target_components if comp != "root") in_root = "root" in target_components and "/" not in file in_target_component = in_target_folder or in_root # Skip huge weights if root processing isn't requested explicitly if "/" not in file and file.endswith(".safetensors") and not in_root: yield f"🗑️ Auto-skipping root model file: {file}..." continue yield f"⏳ Processing {file}..." try: os.makedirs(cache_dir, exist_ok=True) local_path = hf_hub_download(repo_id=source_repo, filename=file, cache_dir=cache_dir, token=token) if file.endswith(".safetensors") and in_target_component: yield f"🧠 Quantizing {file} to {precision}..." tensors = load_file(local_path) new_tensors = {} for k, v in tensors.items(): if not v.is_floating_point(): new_tensors[k] = v continue # Quantization execution logic is_2d_weight = "weight" in k and len(v.shape) == 2 is_excluded = any(ex in k for ex in exclude_prefixes) if is_int8 and is_2d_weight and not is_excluded: scale = v.abs().max(dim=1, keepdim=True)[0] / 127.0 scale = scale.clamp(min=1e-8) new_tensors[f"{k.rsplit('.', 1)[0]}.weight_int8"] = torch.round(v / scale).clamp(-127, 127).to(torch.int8) new_tensors[f"{k.rsplit('.', 1)[0]}.weight_scale"] = scale.to(torch.bfloat16) elif is_int4 and is_2d_weight and not is_excluded: # Standard 4-bit uniform quantization (-8 to 7 range) stored in int8 containers scale = v.abs().max(dim=1, keepdim=True)[0] / 7.0 scale = scale.clamp(min=1e-8) new_tensors[f"{k.rsplit('.', 1)[0]}.weight_int4"] = torch.round(v / scale).clamp(-8, 7).to(torch.int8) new_tensors[f"{k.rsplit('.', 1)[0]}.weight_scale"] = scale.to(torch.bfloat16) elif is_int8 or is_int4: # Fallback for excluded layers or 1D/3D vectors under integer workflows new_tensors[k] = v.to(torch.bfloat16) else: # Casting paths (BF16 / FP16) new_tensors[k] = v.to(target_dtype) converted_path = "converted.safetensors" save_file(new_tensors, converted_path) del tensors, new_tensors gc.collect() yield f"☁️ Uploading processed version of {file}..." api.upload_file(path_or_fileobj=converted_path, path_in_repo=file, repo_id=target_repo) os.remove(converted_path) else: yield f"☁️ Copying non-weight configuration asset: {file}..." api.upload_file(path_or_fileobj=local_path, path_in_repo=file, repo_id=target_repo) success_count += 1 if os.path.exists(cache_dir): shutil.rmtree(cache_dir) gc.collect() except Exception as e: error_count += 1 yield f"⚠️ Error processing {file}: {str(e)}\nSkipping..." if os.path.exists(cache_dir): shutil.rmtree(cache_dir) yield f"✅ Processing Complete! Successfully moved: {success_count} files | Errors: {error_count}." # --- UI Helpers --- def generate_target_repo(source, precision): model_name = source.split("/")[-1] if "/" in source else source return f"your-username/{model_name}-{precision.lower()}" def update_precision_warning(precision): if precision in ["INT8", "INT4"]: return gr.update( value=f"⚠️ **{precision} Warning:** Weight keys will split into binary matrix and scalar scales. Requires custom hardware execution kernels to load.", visible=True ) return gr.update(visible=False) # --- GUI Definition --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown( """ # ❤️ HeartMuLa Dedicated Quantization Suite Optimize, quantize, and shard HeartMuLa architecture weights natively on Hugging Face infrastructure. """ ) with gr.Row(): # Configuration Left Column with gr.Column(scale=6): with gr.Group(): gr.Markdown("### 1. Repository Credentials") hf_token = gr.Textbox(label="Hugging Face Write Token", type="password", placeholder="hf_...") source_repo = gr.Textbox( label="Source Repository Path", placeholder="HeartMuLa/HeartMuLa-oss-3B", info="Target any open HeartMuLa model repository ID." ) gr.Markdown("⚡ **Model Quick Presets**") with gr.Row(): preset_3b = gr.Button("HeartMuLa-3B", size="sm") preset_rl = gr.Button("HeartMuLa-RL", size="sm") preset_codec = gr.Button("HeartCodec", size="sm") preset_gen = gr.Button("HeartMuLaGen", size="sm") preset_trans = gr.Button("Transcriptor", size="sm") with gr.Group(): gr.Markdown("### 2. Hyperparameters & Architectural Logic") arch_profile = gr.Radio( choices=list(ARCH_PROFILES.keys()), value="HeartMuLa / Transformer Core (3B / RL / Gen)", label="Layer Exclusion Mask Profile", info="Protects crucial embeddings and structural norms from severe quantization degradation." ) target_components = gr.CheckboxGroup( choices=["root", "transformer", "vae"], value=["root"], label="Target Sub-Locations", info="HeartMuLa arrays live in the 'root' directory. Adjust only if modifying specific forks." ) with gr.Group(): gr.Markdown("### 3. Execution Target") precision = gr.Dropdown( choices=["BF16", "FP16", "INT8", "INT4"], value="BF16", label="Target Precision Type" ) precision_warning = gr.Markdown(visible=False) target_repo = gr.Textbox(label="Output Destination Repository", placeholder="username/model-bf16") start_btn = gr.Button("🚀 Initialize Serverless Quantization", variant="primary", size="lg") # Output Log Right Column with gr.Column(scale=5): gr.Markdown("### Operational Log Output") output_log = gr.Textbox( label="Terminal Session", lines=28, interactive=False, max_lines=35, autoscroll=True ) # Automated Preset Routing Logic preset_3b.click(lambda: ("HeartMuLa/HeartMuLa-oss-3B", "HeartMuLa / Transformer Core (3B / RL / Gen)", ["root"]), outputs=[source_repo, arch_profile, target_components]) preset_rl.click(lambda: ("HeartMuLa/HeartMuLa-RL-oss-3B-20260123", "HeartMuLa / Transformer Core (3B / RL / Gen)", ["root"]), outputs=[source_repo, arch_profile, target_components]) preset_codec.click(lambda: ("HeartMuLa/HeartCodec-oss-20260123", "HeartMuLa / Flow Matching Codec", ["root"]), outputs=[source_repo, arch_profile, target_components]) preset_gen.click(lambda: ("HeartMuLa/HeartMuLaGen", "HeartMuLa / Transformer Core (3B / RL / Gen)", ["root"]), outputs=[source_repo, arch_profile, target_components]) preset_trans.click(lambda: ("HeartMuLa/HeartTranscriptor-oss", "HeartMuLa / Transformer Core (3B / RL / Gen)", ["root"]), outputs=[source_repo, arch_profile, target_components]) # Dynamics & Event Inversions source_repo.change(fn=generate_target_repo, inputs=[source_repo, precision], outputs=[target_repo]) precision.change(fn=generate_target_repo, inputs=[source_repo, precision], outputs=[target_repo]) precision.change(fn=update_precision_warning, inputs=[precision], outputs=[precision_warning]) start_btn.click( fn=convert_and_upload, inputs=[hf_token, source_repo, target_repo, precision, target_components, arch_profile], outputs=[output_log] ) if __name__ == "__main__": demo.launch()