lalanull/diffusiongemma-custom-handler

1import os2import subprocess3import shutil4from huggingface_hub import hf_hub_download5from typing import Dict, Any6 7class EndpointHandler:8    def __init__(self, path: str):9        # 1. Download GGUF model dynamically10        print("Downloading DiffusionGemma Q4_K_M GGUF model...")11        self.model_path = hf_hub_download(12            repo_id="unsloth/diffusiongemma-26B-A4B-it-GGUF",13            filename="diffusiongemma-26B-A4B-it-Q4_K_M.gguf"14        )15        print(f"Model downloaded to: {self.model_path}")16 17        # 2. Clone and build llama.cpp with PR 24423 (diffusiongemma support)18        print("Cloning llama.cpp...")19        if os.path.exists("/tmp/llama.cpp"):20            shutil.rmtree("/tmp/llama.cpp")21            22        subprocess.run("git clone https://github.com/ggml-org/llama.cpp.git /tmp/llama.cpp", shell=True, check=True)23        24        print("Checking out PR 24423 (diffusiongemma branch)...")25        subprocess.run("cd /tmp/llama.cpp && git fetch origin pull/24423/head:diffusiongemma && git checkout diffusiongemma", shell=True, check=True)26        27        # Check if GPU exists28        has_cuda = os.path.exists("/usr/local/cuda") or subprocess.run("nvidia-smi", shell=True, capture_output=True).returncode == 029        30        print(f"Building llama-diffusion-cli (CUDA={has_cuda})...")31        if has_cuda:32            build_cmd = "cd /tmp/llama.cpp && cmake -B build -DGGML_CUDA=ON && cmake --build build --target llama-diffusion-cli -j 4"33        else:34            build_cmd = "cd /tmp/llama.cpp && cmake -B build && cmake --build build --target llama-diffusion-cli -j 4"35            36        subprocess.run(build_cmd, shell=True, check=True)37        self.binary_path = "/tmp/llama.cpp/build/bin/llama-diffusion-cli"38        print("Successfully compiled llama-diffusion-cli!")39 40    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:41        inputs = data.get("inputs", "")42        if not inputs:43            return {"error": "Missing 'inputs' field in payload"}44            45        # Check if GPU is present to set ngl46        has_cuda = os.path.exists("/usr/local/cuda") or os.path.exists("/dev/nvidia0")47        ngl = 99 if has_cuda else 048        49        # Execute llama-diffusion-cli50        cmd = [self.binary_path, "-m", self.model_path, "-p", inputs, "-ngl", str(ngl), "-n", "300"]51        print(f"Executing: {' '.join(cmd)}")52        53        try:54            res = subprocess.run(cmd, capture_output=True, text=True, timeout=120, check=True)55            return {"generated_text": res.stdout}56        except subprocess.CalledProcessError as e:57            return {"error": f"Execution failed: {e.stderr}\nOutput: {e.stdout}"}58        except Exception as e:59            return {"error": str(e)}60