app.py

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()