app.py

import gradio as gr
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer, AutoModel
import tempfile
from huggingface_hub import HfApi
from huggingface_hub import list_models
from packaging import version
import os
import spaces

def hello(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None) -> str:
    # ^ expect a gr.OAuthProfile object as input to get the user's profile
    # if the user is not logged in, profile will be None
    if profile is None:
        return "Hello !"
    return f"Hello {profile.name} !"

def check_model_exists(oauth_token: gr.OAuthToken | None, username, quantization_type, group_size, model_name, quantized_model_name):
    """Check if a model exists in the user's Hugging Face repository."""
    try:
        models = list_models(author=username, token=oauth_token.token)
        model_names = [model.id for model in models]
        if quantized_model_name : 
            repo_name = f"{username}/{quantized_model_name}"
        else : 
            if quantization_type == "int4_weight_only" : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{quantization_type.lower()}-gs_{group_size}"
            else : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{quantization_type.lower()}"

        if repo_name in model_names:
            return f"Model '{repo_name}' already exists in your repository."
        else:
            return None  # Model does not exist
    except Exception as e:
        return f"Error checking model existence: {str(e)}"

def create_model_card(model_name, quantization_type, group_size):
    model_card = f"""---
base_model:
- {model_name}
---

# {model_name} (Quantized)

## Description
This model is a quantized version of the original model `{model_name}`. It has been quantized using {quantization_type} quantization with torchao.

## Quantization Details
- **Quantization Type**: {quantization_type}
- **Group Size**: {group_size if quantization_type == "int4_weight_only" else None}

## Usage
You can use this model in your applications by loading it directly from the Hugging Face Hub:

```python
from transformers import AutoModel

model = AutoModel.from_pretrained("{model_name}")"""
    
    return model_card

@spaces.GPU
def load_model_gpu(model_name, quantization_config, auth_token) : 
    return AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)

def load_model_cpu(model_name, quantization_config, auth_token) : 
    return AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=quantization_config, use_auth_token=auth_token.token)

def quantize_model(model_name, quantization_type, group_size=128, auth_token=None, username=None, device="cuda"):
    print(f"Quantizing model: {quantization_type}")
    if quantization_type == "int4_weight_only" : 
        quantization_config = TorchAoConfig(quantization_type, group_size=group_size)
    else : 
        quantization_config = TorchAoConfig(quantization_type)
    if device == "cuda" : 
        model = load_model_gpu(model_name, quantization_config=quantization_config, auth_token=auth_token)
    else : 
        model = load_model_cpu(model_name, quantization_config=quantization_config, auth_token=auth_token)

    return model

def save_model(model, model_name, quantization_type, group_size=128, username=None, auth_token=None, quantized_model_name=None):
    print("Saving quantized model")
    with tempfile.TemporaryDirectory() as tmpdirname:


        model.save_pretrained(tmpdirname, safe_serialization=False, use_auth_token=auth_token.token)
        if quantized_model_name : 
            repo_name = f"{username}/{quantized_model_name}"
        else : 
            if quantization_type == "int4_weight_only" : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{quantization_type.lower()}-gs_{group_size}"
            else : 
                repo_name = f"{username}/{model_name.split('/')[-1]}-torchao-{quantization_type.lower()}"

        model_card = create_model_card(repo_name, quantization_type, group_size)
        with open(os.path.join(tmpdirname, "README.md"), "w") as f:
            f.write(model_card)
        # Push to Hub
        api = HfApi(token=auth_token.token)
        api.create_repo(repo_name, exist_ok=True)
        api.upload_folder(
            folder_path=tmpdirname,
            repo_id=repo_name,
            repo_type="model",
        )
    
    return f"https://huggingface.co/{repo_name}"

def quantize_and_save(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None, model_name, quantization_type, group_size, quantized_model_name, device):
    if oauth_token is None : 
        return "Error : Please Sign In to your HuggingFace account to use the quantizer"
    if not profile:
        return "Error: Please Sign In to your HuggingFace account to use the quantizer"
    exists_message = check_model_exists(oauth_token, profile.username, quantization_type, group_size, model_name, quantized_model_name)
    if exists_message : 
        return exists_message
    if quantization_type == "int4_weight_only" and device == "cpu" : 
        return "int4_weight_only not supported on cpu"
    # try : 
    quantized_model = quantize_model(model_name, quantization_type, group_size, oauth_token, profile.username, device)
    return save_model(quantized_model, model_name, quantization_type, group_size, profile.username, oauth_token, quantized_model_name)
    # except Exception as e : 
    #     return e


with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        # 🚀 LLM Model Quantization App
        
        Quantize your favorite Hugging Face models and save them to your profile!
        """
    )

    
    gr.LoginButton(elem_id="login-button", elem_classes="center-button")

    m1 = gr.Markdown()
    app.load(hello, inputs=None, outputs=m1)

    with gr.Row():
        with gr.Column():
            model_name = gr.Textbox(
                label="Model Name",
                placeholder="e.g., meta-llama/Meta-Llama-3-8B",
                value="meta-llama/Meta-Llama-3-8B"
            )
            quantization_type = gr.Dropdown(
                label="Quantization Type",
                choices=["int4_weight_only", "int8_weight_only", "int8_dynamic_activation_int8_weight"],
                value="int8_weight_only"
            )
            group_size = gr.Number(
                label="Group Size (only for int4_weight_only)",
                value=128,
                interactive=True
            )
            device = gr.Dropdown(
                label="Device (int4 only works with cuda)",
                choices=["cuda", "cpu"],
                value="cuda"
            )
            quantized_model_name = gr.Textbox(
                label="Model Name (optional : to override default)",
                value="",
                interactive=True
            )
            # with gr.Row():
            #     username = gr.Textbox(
            #         label="Hugging Face Username",
            #         placeholder="Enter your Hugging Face username",
            #         value="",
            #         interactive=True,
            #         elem_id="username-box"
            #     )
        with gr.Column():
            quantize_button = gr.Button("Quantize and Save Model", variant="primary")
            output_link = gr.Textbox(label="Quantized Model Link")
    
    gr.Markdown(
        """
        ## Instructions
        1. Login to your HuggingFace account
        2. Enter the name of the Hugging Face LLM model you want to quantize (Make sure you have access to it)
        3. Choose the quantization type.
        4. Optionally, specify the group size.
        5. Optionally, choose a custom name for the quantized model
        6. Click "Quantize and Save Model" to start the process.
        7. Once complete, you'll receive a link to the quantized model on Hugging Face.
        
        Note: This process may take some time depending on the model size and your hardware you can check the container logs to see where are you at in the process!
        """
    )
    
    
    # Adding CSS styles for the username box
    app.css = """
    #username-box {
        background-color: #f0f8ff; /* Light color */
        border-radius: 8px;
        padding: 10px;
    }
    """
    app.css = """
    .center-button {
        display: flex;
        justify-content: center;
        align-items: center;
        margin: 0 auto; /* Center horizontally */
    }
    """
    
    quantize_button.click(
        fn=quantize_and_save,
        inputs=[model_name, quantization_type, group_size, quantized_model_name, device],
        outputs=[output_link]
    )


# Launch the app
app.launch()