import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Scalai/scal-lite-60b-code"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

_model_holder = {"model": None}

@spaces.GPU(duration=180)  # generous — first load under MXFP4 can be slow
def _load_model():
    if _model_holder["model"] is None:
        _model_holder["model"] = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype="auto",
            device_map="cuda",
            trust_remote_code=True,
        )
    return _model_holder["model"]

# Force the load to happen under a REAL GPU context, not emulation
_load_model()

@spaces.GPU(duration=90)
def generate(prompt, system_prompt, max_new_tokens, temperature, top_p):
    model = _model_holder["model"]
    # ... rest of your generate function, using `model` instead of the old global