setup.py

"""
setup.py - Run this at the start of app.py to ensure triton_kernels is installed
Add this to the top of your app.py file in HF Spaces
"""

import subprocess
import sys

def ensure_triton_kernels():
    """Ensure triton_kernels is installed for MX format support."""
    try:
        import triton_kernels
        print("✓ triton_kernels already installed")
        return True
    except ImportError:
        print("Installing triton_kernels for MX format support...")
        try:
            subprocess.check_call([
                sys.executable, "-m", "pip", "install",
                "git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels"
            ])
            print("✓ triton_kernels installed successfully")
            return True
        except subprocess.CalledProcessError as e:
            print(f"✗ Failed to install triton_kernels: {e}")
            print("WARNING: MX format will fall back to bf16, LoRA may not work!")
            return False

# Run at import time
if __name__ != "__main__":  # When imported
    ensure_triton_kernels()