app.py

import os
import subprocess
import signal
import tempfile
from pathlib import Path
from textwrap import dedent
import logging
import gradio as gr
from huggingface_hub import HfApi, ModelCard, whoami
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from apscheduler.schedulers.background import BackgroundScheduler
from datetime import datetime
import numpy as np
HF_TOKEN = os.environ.get("HF_TOKEN")

os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
CONVERSION_SCRIPT = "./llama.cpp/convert_hf_to_gguf.py"

log_dir = "/data/logs"
downloads_dir = "/data/downloads"
outputs_dir = "/data/outputs"
os.makedirs(log_dir, exist_ok=True)

logging.basicConfig(
    filename=os.path.join(log_dir, "app.log"),
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

logger = logging.getLogger(__name__)

def get_llama_cpp_notes(gguf_files, new_repo_url, split_model, model_id = None,):
    try:
        result = subprocess.run(
            ['git', '-C', './llama.cpp', 'describe', '--tags', '--always'],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            check=True,
            text=True
        )
        version = result.stdout.strip().split('-')[0]
        text = f"""
*Produced by [Antigma Labs](https://antigma.ai)*
## llama.cpp quantization
Using <a href="https://github.com/ggml-org/llama.cpp">llama.cpp</a> release <a href="https://github.com/ggml-org/llama.cpp/releases/tag/{version}">b4944</a> for quantization.
Original model: https://huggingface.co/{model_id}
Run them directly with [llama.cpp](https://github.com/ggml-org/llama.cpp), or any other llama.cpp based project
## Prompt format
```
<｜begin▁of▁sentence｜>{{system_prompt}}<｜User｜>{{prompt}}<｜Assistant｜><｜end▁of▁sentence｜><｜Assistant｜>
```
## Download a file (not the whole branch) from below:
| Filename | Quant type | File Size | Split |
| -------- | ---------- | --------- | ----- |
| {'|'.join(['|'.join([gguf_files[i][0][:-5] if split_model else ('['+gguf_files[i][0]+']'+'(' + new_repo_url+'/blob/main/'+gguf_files[i][0] + ')'), gguf_files[i][3], f"{gguf_files[i][2]:.2f}" + ' GB', str(split_model),'''
''']) for i in range(len(gguf_files))]) } 
## Downloading using huggingface-cli
<details>
  <summary>Click to view download instructions</summary>
First, make sure you have hugginface-cli installed:
```
pip install -U "huggingface_hub[cli]"
```
Then, you can target the specific file you want:
```
huggingface-cli download {new_repo_url} --include "{gguf_files[0][0]}" --local-dir ./
```
If the model is bigger than 50GB, it will have been split into multiple files. In order to download them all to a local folder, run:
```
huggingface-cli download {new_repo_url} --include "{gguf_files[0][0]}/*" --local-dir ./
```
You can either specify a new local-dir (deepseek-ai_DeepSeek-V3-0324-Q8_0) or download them all in place (./)
</details>
"""
        return text
    except subprocess.CalledProcessError as e:
        print("Error:", e.stderr.strip())
        return None


def get_repo_namespace(repo_owner, username, user_orgs):
    if repo_owner == 'self':
        return username
    for org in user_orgs:
        if org['name'] == repo_owner:
            return org['name']
    raise ValueError(f"Invalid repo_owner: {repo_owner}")

def escape(s: str) -> str:
    return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;").replace("\n", "<br/>")

def toggle_repo_owner(export_to_org, oauth_token: gr.OAuthToken | None):
    if oauth_token is None or oauth_token.token is None:
        raise gr.Error("You must be logged in to use GGUF-my-repo")
    if not export_to_org:
        return gr.update(visible=False, choices=["self"], value="self"), gr.update(visible=False, value="")
    info = whoami(oauth_token.token)
    orgs = [org["name"] for org in info.get("orgs", [])]
    return gr.update(visible=True, choices=["self"] + orgs, value="self"), gr.update(visible=True)
def generate_importance_matrix(model_path: str, train_data_path: str, output_path: str):
    imatrix_command = [
        "./llama.cpp/llama-imatrix",
        "-m", model_path,
        "-f", train_data_path,
        "-ngl", "99",
        "--output-frequency", "10",
        "-o", output_path,
    ]

    if not os.path.isfile(model_path):
        raise Exception(f"Model file not found: {model_path}")

    print("Running imatrix command...")
    process = subprocess.Popen(imatrix_command, shell=False)

    try:
        process.wait(timeout=60)  # added wait
    except subprocess.TimeoutExpired:
        print("Imatrix computation timed out. Sending SIGINT to allow graceful termination...")
        process.send_signal(signal.SIGINT)
        try:
            process.wait(timeout=5)  # grace period
        except subprocess.TimeoutExpired:
            print("Imatrix proc still didn't term. Forecfully terming process...")
            process.kill()

    print("Importance matrix generation completed.")

def split_upload_model(model_path: str, outdir: str, repo_id: str, oauth_token: gr.OAuthToken | None, split_max_tensors=256, split_max_size=None, org_token=None, export_to_org=False):
    print(f"Model path: {model_path}")
    print(f"Output dir: {outdir}")

    if oauth_token is None or oauth_token.token is None:
        raise ValueError("You have to be logged in.")

    split_cmd = [
        "./llama.cpp/llama-gguf-split",
        "--split",
    ]
    if split_max_size:
        split_cmd.append("--split-max-size")
        split_cmd.append(split_max_size)
    else:
        split_cmd.append("--split-max-tensors")
        split_cmd.append(str(split_max_tensors))

    # args for output
    model_path_prefix = '.'.join(model_path.split('.')[:-1]) # remove the file extension
    split_cmd.append(model_path)
    split_cmd.append(model_path_prefix)

    print(f"Split command: {split_cmd}")

    result = subprocess.run(split_cmd, shell=False, capture_output=True, text=True)
    print(f"Split command stdout: {result.stdout}")
    print(f"Split command stderr: {result.stderr}")

    if result.returncode != 0:
        stderr_str = result.stderr.decode("utf-8")
        raise Exception(f"Error splitting the model: {stderr_str}")
    print("Model split successfully!")

    # remove the original model file if needed
    if os.path.exists(model_path):
        os.remove(model_path)

    model_file_prefix = model_path_prefix.split('/')[-1]
    print(f"Model file name prefix: {model_file_prefix}")
    sharded_model_files = [f for f in os.listdir(outdir) if f.startswith(model_file_prefix) and f.endswith(".gguf")]
    if sharded_model_files:
        print(f"Sharded model files: {sharded_model_files}")
        if export_to_org and org_token!="":
          api = HfApi(token = org_token)
        else:
          api = HfApi(token=oauth_token.token)
        for file in sharded_model_files:
            file_path = os.path.join(outdir, file)
            print(f"Uploading file: {file_path}")
            try:
                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=file,
                    repo_id=repo_id,
                )
            except Exception as e:
                raise Exception(f"Error uploading file {file_path}: {e}")
    else:
        raise Exception("No sharded files found.")

    print("Sharded model has been uploaded successfully!")

def process_model(model_id, q_method, use_imatrix, imatrix_q_method, private_repo,
                  train_data_file, split_model, split_max_tensors, split_max_size,
                  export_to_org, repo_owner, org_token, oauth_token: gr.OAuthToken | None):
    if oauth_token is None or oauth_token.token is None:
        raise gr.Error("You must be logged in to use GGUF-my-repo")

    user_info = whoami(oauth_token.token)
    username = user_info["name"]
    user_orgs = user_info.get("orgs", [])
    if not export_to_org:
        repo_owner = "self"


    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    logger.info(f"Time {current_time}, Username {username}, Model_ID, {model_id}, q_method {','.join(q_method)}")

    repo_namespace = get_repo_namespace(repo_owner, username, user_orgs)
    model_name = model_id.split('/')[-1]
    try:
      api_token = org_token if (export_to_org and org_token!="") else oauth_token.token
      api = HfApi(token=api_token)

      dl_pattern = ["*.md", "*.json", "*.model"]
      pattern = "*.safetensors" if any(
          f.path.endswith(".safetensors")
          for f in api.list_repo_tree(repo_id=model_id, recursive=True)
      ) else "*.bin"
      dl_pattern += [pattern]

      os.makedirs(downloads_dir, exist_ok=True)
      os.makedirs(outputs_dir, exist_ok=True)

      with tempfile.TemporaryDirectory(dir=outputs_dir) as outdir:
          fp16 = str(Path(outdir)/f"{model_name}.fp16.gguf")

          with tempfile.TemporaryDirectory(dir=downloads_dir) as tmpdir:
              print("Downloading")
              local_dir = Path(tmpdir)/model_name
              api.snapshot_download(repo_id=model_id, local_dir=local_dir, local_dir_use_symlinks=False, allow_patterns=dl_pattern)
              
              config_dir = local_dir/"config.json"
              adapter_config_dir = local_dir/"adapter_config.json"
              if os.path.exists(adapter_config_dir) and not os.path.exists(config_dir):
                  raise Exception("adapter_config.json is present. If converting LoRA, use GGUF-my-lora.")

              print("Download successfully")
              result = subprocess.run(["python", CONVERSION_SCRIPT, local_dir, "--outtype", "f16", "--outfile", fp16], shell=False, capture_output=True)
              print("Converted to f16")
              if result.returncode != 0:
                  raise Exception(f"Error converting to fp16: {result.stderr.decode()}")

          imatrix_path = Path(outdir)/"imatrix.dat"
          if use_imatrix:
              train_data_path = train_data_file.name if train_data_file else "llama.cpp/groups_merged.txt"
              if not os.path.isfile(train_data_path):
                  raise Exception(f"Training data not found: {train_data_path}")
              generate_importance_matrix(fp16, train_data_path, imatrix_path)

          quant_methods = [imatrix_q_method] if use_imatrix else (q_method if isinstance(q_method, list) else [q_method])
          suffix = "imat" if use_imatrix else None

          gguf_files = []
          for method in quant_methods:
              print("Begin quantize")
              name = f"{model_name.lower()}-{method.lower()}-{suffix}.gguf" if suffix else f"{model_name.lower()}-{method.lower()}.gguf"
              path = str(Path(outdir)/name)
              quant_cmd = ["./llama.cpp/llama-quantize", "--imatrix", imatrix_path, fp16, path, method] if use_imatrix else ["./llama.cpp/llama-quantize", fp16, path, method]
              result = subprocess.run(quant_cmd, shell=False, capture_output=True)
              if result.returncode != 0:
                  raise Exception(f"Quantization failed ({method}): {result.stderr.decode()}")
              size = os.path.getsize(path)/1024/1024/1024
              gguf_files.append((name, path, size, method))

          print("Quantize successfully!")
          suffix_for_repo = f"{imatrix_q_method}-imat" if use_imatrix else "-".join(quant_methods)
          repo_id = f"{repo_namespace}/{model_name}-{suffix_for_repo}-GGUF"
          new_repo_url = api.create_repo(repo_id=repo_id, exist_ok=True, private=private_repo)

          try:
              card = ModelCard.load(model_id, token=oauth_token.token)
          except:
              card = ModelCard("")
          card.data.tags = (card.data.tags or []) + ["llama-cpp", "gguf-my-repo"]
          card.data.base_model = model_id
          card.text = dedent(get_llama_cpp_notes(gguf_files, new_repo_url, split_model, model_id))
          readme_path = Path(outdir)/"README.md"
          card.save(readme_path)
          for name, path, _, _ in gguf_files:
              if split_model:
                  split_upload_model(path, outdir, repo_id, oauth_token, split_max_tensors, split_max_size, org_token, export_to_org)
              else:
                  api.upload_file(path_or_fileobj=path, path_in_repo=name, repo_id=repo_id)
          if use_imatrix and os.path.isfile(imatrix_path):
              api.upload_file(path_or_fileobj=imatrix_path, path_in_repo="imatrix.dat", repo_id=repo_id)
          api.upload_file(path_or_fileobj=readme_path, path_in_repo="README.md", repo_id=repo_id)

          return (f'<h1>✅ DONE</h1><br/>Repo: <a href="{new_repo_url}" target="_blank" style="text-decoration:underline">{repo_id}</a>', f"llama{np.random.randint(9)}.png")
    except Exception as e:
        raise (f'<h1>❌ ERROR</h1><br/><pre style="white-space:pre-wrap;">{escape(str(e))}</pre>', "error.png")


css="""/* Custom CSS to allow scrolling */
.gradio-container {overflow-y: auto;}
"""
model_id = HuggingfaceHubSearch(
    label="Hub Model ID",
    placeholder="Search for model id on Huggingface",
    search_type="model",
)

export_to_org = gr.Checkbox(
    label="Export to Organization Repository",
    value=False,
    info="If checked, you can select an organization to export to."
)

repo_owner = gr.Dropdown(
    choices=["self"],
    value="self",
    label="Repository Owner",
    visible=False
)

org_token = gr.Textbox(
    label="Org Access Token",
    type="password",
    visible=False
)

q_method = gr.Dropdown(
    ["Q2_K", "Q3_K_S", "Q3_K_M", "Q3_K_L", "Q4_0", "Q4_K_S", "Q4_K_M", "Q5_0", "Q5_K_S", "Q5_K_M", "Q6_K", "Q8_0"],
    label="Quantization Method",
    info="GGML quantization type",
    value="Q4_K_M",
    filterable=False,
    visible=True,
    multiselect=True
)

imatrix_q_method = gr.Dropdown(
    ["IQ3_M", "IQ3_XXS", "Q4_K_M", "Q4_K_S", "IQ4_NL", "IQ4_XS", "Q5_K_M", "Q5_K_S"],
    label="Imatrix Quantization Method",
    info="GGML imatrix quants type",
    value="IQ4_NL",
    filterable=False,
    visible=False
)

use_imatrix = gr.Checkbox(
    value=False,
    label="Use Imatrix Quantization",
    info="Use importance matrix for quantization."
)

private_repo = gr.Checkbox(
    value=False,
    label="Private Repo",
    info="Create a private repo under your username."
)

train_data_file = gr.File(
    label="Training Data File",
    file_types=["txt"],
    visible=False
)

split_model = gr.Checkbox(
    value=False,
    label="Split Model",
    info="Shard the model using gguf-split."
)

split_max_tensors = gr.Number(
    value=256,
    label="Max Tensors per File",
    info="Maximum number of tensors per file when splitting model.",
    visible=False
)

split_max_size = gr.Textbox(
    label="Max File Size",
    info="Maximum file size when splitting model (--split-max-size). May leave empty to use the default. Accepted suffixes: M, G. Example: 256M, 5G",
    visible=False
)

iface = gr.Interface(
    fn=process_model,
    inputs=[
        model_id,
        q_method,
        use_imatrix,
        imatrix_q_method,
        private_repo,
        train_data_file,
        split_model,
        split_max_tensors,
        split_max_size,
        export_to_org,
        repo_owner,
        org_token
    ],
    outputs=[
        gr.Markdown(label="Output"),
        gr.Image(show_label=False)
    ],
    title="Make your own GGUF Quants — faster than ever before, believe me.",
    description="We take your Hugging Face repo — a terrific repo — we quantize it, we package it beautifully, and we give you your very own repo. It's smart. It's efficient. It's huge. You're gonna love it.",
    api_name=False
)
with gr.Blocks(css=".gradio-container {overflow-y: auto;}") as demo:
    gr.Markdown("Logged in, you must be. Classy, secure, and victorious, it keeps us.")
    gr.LoginButton(min_width=250)



    export_to_org.change(fn=toggle_repo_owner, inputs=[export_to_org], outputs=[repo_owner, org_token])

    split_model.change(fn=lambda sm: (gr.update(visible=sm), gr.update(visible=sm)), inputs=split_model, outputs=[split_max_tensors, split_max_size])
    use_imatrix.change(fn=lambda use: (gr.update(visible=not use), gr.update(visible=use), gr.update(visible=use)), inputs=use_imatrix, outputs=[q_method, imatrix_q_method, train_data_file])

    iface.render()


def restart_space():
    HfApi().restart_space(repo_id="Antigma/quantize-my-repo", token=HF_TOKEN, factory_reboot=True)

scheduler = BackgroundScheduler()
scheduler.add_job(restart_space, "interval", seconds=21600)
scheduler.start()

demo.queue(default_concurrency_limit=1, max_size=5).launch(debug=True, show_api=False)