app.py

import gradio as gr
import torch
from transformers import AutoModel, BitsAndBytesConfig
import tempfile
from huggingface_hub import HfApi
from huggingface_hub import list_models
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from bitsandbytes.nn import Linear4bit
from packaging import version
import os


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 Please Login to HuggingFace to use the BitsAndBytes Quantizer!"
    return f"Hello {profile.name} ! Welcome to BitsAndBytes Quantizer"


def check_model_exists(
    oauth_token: gr.OAuthToken | None, username, 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:
            repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"

        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, quant_type_4, double_quant_4, compute_type_4, quant_storage_4
):
    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 int4 quantization with bitsandbytes.

## Quantization Details
- **Quantization Type**: int4
- **bnb_4bit_quant_type**: {quant_type_4}
- **bnb_4bit_use_double_quant**: {double_quant_4}
- **bnb_4bit_compute_dtype**: {compute_type_4}
- **bnb_4bit_quant_storage**: {quant_storage_4}

## 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


DTYPE_MAPPING = {
    "int8": torch.int8,
    "uint8": torch.uint8,
    "float16": torch.float16,
    "float32": torch.float32,
    "bfloat16": torch.bfloat16,
}


def quantize_model(
    model_name,
    quant_type_4,
    double_quant_4,
    compute_type_4,
    quant_storage_4,
    auth_token=None,
    progress=gr.Progress(),
):
    progress(0, desc="Loading model")

    # Configure quantization
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type=quant_type_4,
        bnb_4bit_use_double_quant=True if double_quant_4 == "True" else False,
        bnb_4bit_quant_storage=DTYPE_MAPPING[quant_storage_4],
        bnb_4bit_compute_dtype=DTYPE_MAPPING[compute_type_4],
    )

    # Load model
    model = AutoModel.from_pretrained(
        model_name,
        quantization_config=quantization_config,
        device_map="cpu",
        use_auth_token=auth_token.token,
        torch_dtype=torch.bfloat16,
    )
    progress(0.33, desc="Quantizing")

    # Quantize model
    modules = list(model.named_modules())
    for idx, (_, module) in enumerate(modules):
        if isinstance(module, Linear4bit):
            module.to("cuda")
            module.to("cpu")
        progress(0.33 + (0.33 * idx / len(modules)), desc="Quantizing")

    progress(0.66, desc="Quantized successfully")
    return model


def save_model(
    model,
    model_name,
    quant_type_4,
    double_quant_4,
    compute_type_4,
    quant_storage_4,
    username=None,
    auth_token=None,
    quantized_model_name=None,
    public=False,
    progress=gr.Progress(),
):
    progress(0.67, desc="Preparing to push")

    with tempfile.TemporaryDirectory() as tmpdirname:
        # Save model
        model.save_pretrained(
            tmpdirname, safe_serialization=True, use_auth_token=auth_token.token
        )
        progress(0.75, desc="Preparing to push")

        # Prepare repo name and model card
        if quantized_model_name:
            repo_name = f"{username}/{quantized_model_name}"
        else:
            repo_name = f"{username}/{model_name.split('/')[-1]}-bnb-4bit"

        model_card = create_model_card(
            repo_name, quant_type_4, double_quant_4, compute_type_4, quant_storage_4
        )
        with open(os.path.join(tmpdirname, "README.md"), "w") as f:
            f.write(model_card)
        progress(0.80, desc="Model card created")

        # Push to Hub
        api = HfApi(token=auth_token.token)
        api.create_repo(repo_name, exist_ok=True, private=not public)
        progress(0.85, desc="Pushing to Hub")

        # Upload files
        api.upload_folder(
            folder_path=tmpdirname,
            repo_id=repo_name,
            repo_type="model",
        )
        progress(1.00, desc="Model pushed to Hub")

    # Get model architecture as string
    import io
    from contextlib import redirect_stdout
    import html

    # Capture the model architecture string
    f = io.StringIO()
    with redirect_stdout(f):
        print(model)
    model_architecture_str = f.getvalue()

    # Escape HTML characters and format with line breaks
    model_architecture_str_html = html.escape(model_architecture_str).replace(
        "\n", "<br/>"
    )

    # Format it for display in markdown with proper styling
    model_architecture_info = f"""
    <div class="model-architecture" style="max-height: 500px; overflow-y: auto; overflow-x: auto; background-color: #f5f5f5; padding: 5px; border-radius: 8px; font-family: monospace; white-space: pre-wrap;">
        <div style="line-height: 1.2; font-size: 0.75em;">{model_architecture_str_html}</div>
    </div>
    """
    return f'🔗 Quantized Model <br/><h1> 🤗 DONE</h1><br/>Find your repo here: <a href="https://huggingface.co/{repo_name}" target="_blank" style="text-decoration:underline">{repo_name}</a><br/><br/>📊 Model Architecture<br/>{model_architecture_info}'


def quantize_and_save(
    profile: gr.OAuthProfile | None,
    oauth_token: gr.OAuthToken | None,
    model_name,
    quant_type_4,
    double_quant_4,
    compute_type_4,
    quant_storage_4,
    quantized_model_name,
    public,
    progress=gr.Progress(),
):
    if oauth_token is None:
        return """
        <div class="error-box">
            <h3>❌ Authentication Error</h3>
            <p>Please sign in to your HuggingFace account to use the quantizer.</p>
        </div>
        """
    if not profile:
        return """
        <div class="error-box">
            <h3>❌ Authentication Error</h3>
            <p>Please sign in to your HuggingFace account to use the quantizer.</p>
        </div>
        """
    exists_message = check_model_exists(
        oauth_token, profile.username, model_name, quantized_model_name
    )
    if exists_message:
        return f"""
        <div class="warning-box">
            <h3>⚠️ Model Already Exists</h3>
            <p>{exists_message}</p>
        </div>
        """
    try:
        # Download phase
        progress(0, desc="Starting quantization process")
        quantized_model = quantize_model(
            model_name,
            quant_type_4,
            double_quant_4,
            compute_type_4,
            quant_storage_4,
            oauth_token,
            progress,
        )
        final_message = save_model(
            quantized_model,
            model_name,
            quant_type_4,
            double_quant_4,
            compute_type_4,
            quant_storage_4,
            profile.username,
            oauth_token,
            quantized_model_name,
            public,
            progress,
        )

        return final_message

    except Exception as e:
        error_message = str(e).replace("\n", "<br/>")
        return f"""
        <div class="error-box">
            <h3>❌ Error Occurred</h3>
            <p>{error_message}</p>
        </div>
        """


css = """/* Custom CSS to allow scrolling */
.gradio-container {overflow-y: auto;}

/* Fix alignment for radio buttons and checkboxes */
.gradio-radio {
    display: flex !important;
    align-items: center !important;
    margin: 10px 0 !important;
}

.gradio-checkbox {
    display: flex !important;
    align-items: center !important;
    margin: 10px 0 !important;
}

/* Ensure consistent spacing and alignment */
.gradio-dropdown, .gradio-textbox, .gradio-radio, .gradio-checkbox {
    margin-bottom: 12px !important;
    width: 100% !important;
}

/* Align radio buttons and checkboxes horizontally */
.option-row {
    display: flex !important;
    justify-content: space-between !important;
    align-items: center !important;
    gap: 20px !important;
    margin-bottom: 12px !important;
}

.option-row .gradio-radio, .option-row .gradio-checkbox {
    margin: 0 !important;
    flex: 1 !important;
}

/* Horizontally align radio button options with text */
.gradio-radio label {
    display: flex !important;
    align-items: center !important;
}

.gradio-radio input[type="radio"] {
    margin-right: 5px !important;
}

/* Remove padding and margin from model name textbox for better alignment */
.model-name-textbox {
    padding-left: 0 !important;
    padding-right: 0 !important;
    margin-left: 0 !important;
    margin-right: 0 !important;
}

/* Quantize button styling with glow effect */
button[variant="primary"] {
    background: linear-gradient(135deg, #3B82F6, #10B981) !important;
    color: white !important;
    padding: 16px 32px !important;
    font-size: 1.1rem !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 12px !important;
    box-shadow: 0 0 15px rgba(59, 130, 246, 0.5) !important;
    transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important;
    position: relative;
    overflow: hidden;
    animation: glow 1.5s ease-in-out infinite alternate;
}

button[variant="primary"]::before {
    content: "✨ ";
}

button[variant="primary"]:hover {
    transform: translateY(-5px) scale(1.05) !important;
    box-shadow: 0 10px 25px rgba(59, 130, 246, 0.7) !important;
}

@keyframes glow {
    from {
        box-shadow: 0 0 10px rgba(59, 130, 246, 0.5);
    }
    to {
        box-shadow: 0 0 20px rgba(59, 130, 246, 0.8), 0 0 30px rgba(16, 185, 129, 0.5);
    }
}

/* Login button styling with glow effect */
#login-button {
    background: linear-gradient(135deg, #3B82F6, #10B981) !important;
    color: white !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 12px !important;
    box-shadow: 0 0 15px rgba(59, 130, 246, 0.5) !important;
    transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important;
    position: relative;
    overflow: hidden;
    animation: glow 1.5s ease-in-out infinite alternate;
    max-width: 300px !important;
    margin: 0 auto !important;
}

#login-button::before {
    content: "🔑 ";
    display: inline-block !important;
    vertical-align: middle !important;
    margin-right: 5px !important;
    line-height: normal !important;
}

#login-button:hover {
    transform: translateY(-3px) scale(1.03) !important;
    box-shadow: 0 10px 25px rgba(59, 130, 246, 0.7) !important;
}

#login-button::after {
    content: "";
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.2), transparent);
    transition: 0.5s;
}

#login-button:hover::after {
    left: 100%;
}

/* Toggle instructions button styling */
#toggle-button {
    background: linear-gradient(135deg, #3B82F6, #10B981) !important;
    color: white !important;
    font-size: 0.85rem !important;
    font-weight: 600 !important;
    padding: 8px 16px !important;
    border: none !important;
    border-radius: 8px !important;
    box-shadow: 0 2px 10px rgba(59, 130, 246, 0.3) !important;
    transition: all 0.3s ease !important;
    margin: 0.5rem auto 1.5rem auto !important;
    display: block !important;
    max-width: 200px !important;
    text-align: center !important;
    position: relative;
    overflow: hidden;
}

#toggle-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 12px rgba(59, 130, 246, 0.5) !important;
}

#toggle-button::after {
    content: "";
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.2), transparent);
    transition: 0.5s;
}

#toggle-button:hover::after {
    left: 100%;
}
/* Progress Bar Styles */
.progress-container {
    font-family: system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
    padding: 20px;
    background: white;
    border-radius: 12px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}

.progress-stage {
    font-size: 0.9rem;
    font-weight: 600;
    color: #64748b;
}

.progress-stage .stage {
    position: relative;
    padding: 8px 12px;
    border-radius: 6px;
    background: #f1f5f9;
    transition: all 0.3s ease;
}

.progress-stage .stage.completed {
    background: #ecfdf5;
}

.progress-bar {
    box-shadow: inset 0 2px 4px rgba(0, 0, 0, 0.1);
}
.progress {
    transition: width 0.8s cubic-bezier(0.4, 0, 0.2, 1);
    box-shadow: 0 2px 4px rgba(59, 130, 246, 0.3);
}
"""


with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
    gr.Markdown(
        """
        # 🤗 LLM Model BitsAndBytes Quantizer ✨
        
        """
    )

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

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

    instructions_visible = gr.State(False)

    with gr.Row():
        with gr.Column():
            with gr.Row():
                model_name = HuggingfaceHubSearch(
                    label="🔍 Hub Model ID",
                    placeholder="Search for model id on Huggingface",
                    search_type="model",
                )
            with gr.Row():
                with gr.Column():
                    gr.Markdown(
                        """
                        ### ⚙️ Model Quantization Type Settings
                        """
                    )
                    quant_type_4 = gr.Dropdown(
                        info="The quantization data type in the bnb.nn.Linear4Bit layers",
                        choices=["fp4", "nf4"],
                        value="nf4",
                        visible=True,
                        show_label=False,
                    )
                    compute_type_4 = gr.Dropdown(
                        info="The compute type for the model",
                        choices=["float16", "bfloat16", "float32"],
                        value="bfloat16",
                        visible=True,
                        show_label=False,
                    )
                    quant_storage_4 = gr.Dropdown(
                        info="The storage type for the model",
                        choices=["float16", "float32", "int8", "uint8", "bfloat16"],
                        value="uint8",
                        visible=True,
                        show_label=False,
                    )
                    gr.Markdown(
                        """
                        ### 🔄 Double Quantization Settings
                        """
                    )
                    with gr.Row(elem_classes="option-row"):
                        double_quant_4 = gr.Radio(
                            ["True", "False"],
                            info="Use Double Quant",
                            visible=True,
                            value="True",
                            show_label=False,
                        )
                    gr.Markdown(
                        """
                        ### 💾 Saving Settings
                        """
                    )
                    with gr.Row():
                        quantized_model_name = gr.Textbox(
                            label="✏️ Model Name",
                            info="Model Name (optional : to override default)",
                            value="",
                            interactive=True,
                            elem_classes="model-name-textbox",
                            show_label=False,
                        )

                    with gr.Row():
                        public = gr.Checkbox(
                            label="🌐 Make model public",
                            info="If checked, the model will be publicly accessible",
                            value=True,
                            interactive=True,
                            show_label=True,
                        )

        with gr.Column():
            quantize_button = gr.Button(
                "🚀 Quantize and Push to the Hub", variant="primary"
            )
            output_link = gr.Markdown(
                "🔗 Quantized Model", container=True, min_height=100
            )

    quantize_button.click(
        fn=quantize_and_save,
        inputs=[
            model_name,
            quant_type_4,
            double_quant_4,
            compute_type_4,
            quant_storage_4,
            quantized_model_name,
            public,
        ],
        outputs=[output_link],
        show_progress="full",
    )
    # Add information section about the app options
    with gr.Accordion("📚 About this app", open=True):
        gr.Markdown(
            """
            ## 📝 Notes on Quantization Options
            
            ### Quantization Type (bnb_4bit_quant_type)
            - **fp4**: Floating-point 4-bit quantization.
            - **nf4**: Normal float 4-bit quantization.
            
            ### Double Quantization
            - **True**: Applies a second round of quantization to the quantization constants, further reducing memory usage.
            - **False**: Uses standard quantization only.
            
            ### Model Saving Options
            - **Model Name**: Custom name for your quantized model on the Hub. If left empty, a default name will be generated.
            - **Make model public**: If checked, anyone can access your quantized model. If unchecked, only you can access it.
            
            ## 🔍 How It Works
            This app uses the BitsAndBytes library to perform 4-bit quantization on Transformer models. The process:
            1. Downloads the original model
            2. Applies the selected quantization settings
            3. Uploads the quantized model to your HuggingFace account
            
            ## 📊 Memory Usage
            4-bit quantization can reduce model size by up to 75% compared to FP16, allowing you to run larger models on consumer hardware.
            """
        )

if __name__ == "__main__":
    demo.launch(share=True)