Deploy Nomos ZeroGPU app

README.md

@@ -17,8 +17,15 @@ This Space runs Nomos-compatible models with ZeroGPU and tries model candidates
 
 ## Suggested Variables
 
-- `MODEL_CANDIDATES=cyankiwi/nomos-1-AWQ-8bit,cyankiwi/nomos-1-AWQ-4bit`
+- `MODEL_CANDIDATES=cyankiwi/nomos-1-AWQ-8bit`
+- `TOKENIZER_ID=NousResearch/nomos-1`
+- `TORCH_DTYPE=bfloat16`
+- `MODEL_DEVICE_MAP=auto`
 - `PREFER_FULL=false`
-- `GPU_DURATION_SECONDS=120`
+- `GPU_SIZE=xlarge`
+- `GPU_DURATION_SECONDS=180`
+- `MAX_GPU_DURATION_SECONDS=300`
 - `MAX_INPUT_TOKENS=2048`
 - `MAX_NEW_TOKENS_DEFAULT=256`
+- `HF_HOME=/tmp/hf-home`
+- `HF_HUB_CACHE=/tmp/hf-home/hub`

app.py

@@ -1,35 +1,41 @@
 #!/usr/bin/env python3
 import os
 import threading
-from collections.abc import Mapping
 from typing import Any
 
-import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
+# Importing spaces early is recommended for ZeroGPU runtime patching.
 try:
     import spaces
 except Exception:
     class _SpacesFallback:
         @staticmethod
-        def GPU(duration: int = 60):
+        def GPU(*args, **kwargs):
             def _decorator(fn):
                 return fn
+
             return _decorator
+
     spaces = _SpacesFallback()
 
+import gradio as gr
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 DEFAULT_FULL_MODEL = "NousResearch/nomos-1"
-DEFAULT_MODEL_CANDIDATES = "cyankiwi/nomos-1-AWQ-8bit,cyankiwi/nomos-1-AWQ-4bit"
+DEFAULT_MODEL_CANDIDATES = "cyankiwi/nomos-1-AWQ-8bit"
 DEFAULT_TOKENIZER_ID = DEFAULT_FULL_MODEL
-GPU_DURATION_SECONDS = int(os.getenv("GPU_DURATION_SECONDS", "120"))
+
+GPU_DURATION_SECONDS = int(os.getenv("GPU_DURATION_SECONDS", "180"))
+MAX_GPU_DURATION_SECONDS = int(os.getenv("MAX_GPU_DURATION_SECONDS", "300"))
+GPU_SIZE = os.getenv("GPU_SIZE", "large").strip().lower() or "large"
 MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "2048"))
 MAX_NEW_TOKENS_DEFAULT = int(os.getenv("MAX_NEW_TOKENS_DEFAULT", "256"))
 TRUST_REMOTE_CODE = os.getenv("TRUST_REMOTE_CODE", "true").lower() == "true"
 PREFER_FULL = os.getenv("PREFER_FULL", "false").lower() == "true"
 TOKENIZER_ID = os.getenv("TOKENIZER_ID", DEFAULT_TOKENIZER_ID).strip() or DEFAULT_TOKENIZER_ID
-TORCH_DTYPE = os.getenv("TORCH_DTYPE", "float16").strip().lower()
+TORCH_DTYPE = os.getenv("TORCH_DTYPE", "bfloat16").strip().lower()
+MODEL_DEVICE_MAP = os.getenv("MODEL_DEVICE_MAP", "auto").strip() or "auto"
 
 _MODEL_LOCK = threading.Lock()
 _MODEL: Any = None
@@ -46,6 +52,45 @@ def _ordered_candidates() -> list[str]:
     return candidates
 
 
+def _torch_dtype() -> torch.dtype | str:
+    if TORCH_DTYPE in {"", "auto"}:
+        return "auto"
+    if TORCH_DTYPE in {"bfloat16", "bf16"}:
+        return torch.bfloat16
+    if TORCH_DTYPE in {"float16", "fp16", "half"}:
+        return torch.float16
+    if TORCH_DTYPE in {"float32", "fp32"}:
+        return torch.float32
+    return "auto"
+
+
+def _package_versions() -> str:
+    pieces = [
+        f"torch={torch.__version__}",
+        f"transformers={transformers.__version__}",
+    ]
+    try:
+        import compressed_tensors
+
+        pieces.append(f"compressed-tensors={compressed_tensors.__version__}")
+    except Exception as exc:  # pragma: no cover - environment specific
+        pieces.append(f"compressed-tensors=unavailable({type(exc).__name__})")
+    return ", ".join(pieces)
+
+
+def _cuda_status() -> str:
+    if not torch.cuda.is_available():
+        return "CUDA unavailable"
+
+    try:
+        idx = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(idx)
+        total_gb = props.total_memory / (1024**3)
+        return f"{props.name} ({total_gb:.1f} GB)"
+    except Exception as exc:  # pragma: no cover - environment specific
+        return f"CUDA available (details unavailable: {type(exc).__name__})"
+
+
 def _load_model_if_needed() -> tuple[str | None, str]:
     global _MODEL, _TOKENIZER, _MODEL_ID
     if _MODEL is not None and _TOKENIZER is not None and _MODEL_ID is not None:
@@ -62,14 +107,21 @@ def _load_model_if_needed() -> tuple[str | None, str]:
                     TOKENIZER_ID,
                     trust_remote_code=TRUST_REMOTE_CODE,
                 )
-                model = AutoModelForCausalLM.from_pretrained(
-                    candidate,
-                    device_map="auto",
-                    trust_remote_code=TRUST_REMOTE_CODE,
-                    low_cpu_mem_usage=True,
-                    torch_dtype=torch.float16 if TORCH_DTYPE == "float16" else torch.bfloat16,
-                )
+                if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+                    tokenizer.pad_token = tokenizer.eos_token
+
+                dtype = _torch_dtype()
+                model_kwargs: dict[str, Any] = {
+                    "trust_remote_code": TRUST_REMOTE_CODE,
+                    "low_cpu_mem_usage": True,
+                    "device_map": MODEL_DEVICE_MAP,
+                }
+                if dtype != "auto":
+                    model_kwargs["torch_dtype"] = dtype
+
+                model = AutoModelForCausalLM.from_pretrained(candidate, **model_kwargs)
                 model.eval()
+
                 _TOKENIZER = tokenizer
                 _MODEL = model
                 _MODEL_ID = candidate
@@ -88,10 +140,12 @@ def _status_text() -> str:
     base = (
         f"Loaded model: `{loaded}`\n\n"
         f"Tokenizer: `{TOKENIZER_ID}`\n\n"
-        f"Torch dtype: `{TORCH_DTYPE}`\n\n"
+        f"Torch dtype: `{TORCH_DTYPE}` | Device map: `{MODEL_DEVICE_MAP}`\n\n"
+        f"GPU size: `{GPU_SIZE}` | Duration default: `{GPU_DURATION_SECONDS}s`\n\n"
+        f"Max input tokens: `{MAX_INPUT_TOKENS}`\n\n"
         f"Candidates: `{candidates}`\n\n"
-        f"GPU duration: `{GPU_DURATION_SECONDS}s` | "
-        f"Max input tokens: `{MAX_INPUT_TOKENS}`"
+        f"Runtime: `{_cuda_status()}`\n\n"
+        f"Packages: `{_package_versions()}`"
     )
     if _LOAD_ERRORS:
         err = "\n".join(f"- {e}" for e in _LOAD_ERRORS[-3:])
@@ -99,7 +153,32 @@ def _status_text() -> str:
     return base
 
 
-@spaces.GPU(duration=GPU_DURATION_SECONDS)
+def _duration_for_generate(
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    do_sample: bool,
+) -> int:
+    del prompt, temperature, top_p, top_k, do_sample
+    try:
+        requested_new = int(max_new_tokens)
+    except Exception:
+        requested_new = MAX_NEW_TOKENS_DEFAULT
+
+    est = max(GPU_DURATION_SECONDS, 60 + int(0.8 * max(32, requested_new)))
+    return min(MAX_GPU_DURATION_SECONDS, est)
+
+
+def _gpu_decorator():
+    try:
+        return spaces.GPU(duration=_duration_for_generate, size=GPU_SIZE)
+    except TypeError:
+        return spaces.GPU(duration=_duration_for_generate)
+
+
+@_gpu_decorator()
 def generate(
     prompt: str,
     max_new_tokens: int,
@@ -120,24 +199,17 @@ def generate(
     model = _MODEL
 
     messages = [{"role": "user", "content": prompt}]
-    chat_inputs = tokenizer.apply_chat_template(
+    chat_text = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
-        return_tensors="pt",
+        tokenize=False,
     )
+    model_inputs = tokenizer(chat_text, return_tensors="pt")
+
     try:
         device = next(model.parameters()).device
     except Exception:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    if hasattr(chat_inputs, "to"):
-        chat_inputs = chat_inputs.to(device)
-
-    if torch.is_tensor(chat_inputs):
-        model_inputs: dict[str, Any] = {"input_ids": chat_inputs}
-    elif isinstance(chat_inputs, Mapping):
-        model_inputs = dict(chat_inputs)
-    else:
-        raise TypeError(f"Unsupported chat template output type: {type(chat_inputs)}")
 
     for k, v in list(model_inputs.items()):
         if torch.is_tensor(v):
@@ -151,12 +223,20 @@ def generate(
                 model_inputs[k] = v[:, trim:]
         input_ids = model_inputs["input_ids"]
 
+    generation_cfg = getattr(model, "generation_config", None)
+    eos_token_id = getattr(generation_cfg, "eos_token_id", None)
+    pad_token_id = getattr(generation_cfg, "pad_token_id", None)
+    if pad_token_id is None:
+        pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id or 0
+
     gen_kwargs: dict[str, Any] = {
         **model_inputs,
         "max_new_tokens": int(max_new_tokens),
         "do_sample": bool(do_sample),
-        "pad_token_id": tokenizer.eos_token_id or tokenizer.pad_token_id or 0,
+        "pad_token_id": pad_token_id,
     }
+    if eos_token_id is not None:
+        gen_kwargs["eos_token_id"] = eos_token_id
     if do_sample:
         gen_kwargs.update(
             {

requirements.txt

@@ -1,6 +1,6 @@
-gradio>=5.0.0
-spaces>=0.30.0
-transformers>=4.51.0
+gradio==5.12.0
+spaces>=0.32.0
+transformers==4.57.3
 accelerate>=0.34.0
 safetensors>=0.5.0
-compressed-tensors>=0.12.3
+compressed-tensors==0.12.3a20251110