v3.4 safe mode: lazy import transformers only when GPU available

app.py

@@ -1,59 +1,50 @@
-import os, spaces, torch, gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+import os, json
+import gradio as gr
 
-SPECS = {
-  'qwen2.5-0.5b': ('Qwen/Qwen2.5-0.5B-Instruct', 1.0, 'public'),
-  'llama3.2-1b': ('meta-llama/Llama-3.2-1B-Instruct', 1.5, 'public'),
-  'deepseek-r1-1.5b': ('deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B', 2.2, 'public'),
-  'gemma3-4b': ('google/gemma-3-4b-it', 6.0, 'public'),
-}
+SAFE = {"qwen2.5-0.5b": ("Qwen/Qwen2.5-0.5B-Instruct", 1.0), "llama3.2-1b": ("meta-llama/Llama-3.2-1B-Instruct", 1.5)}
+SPECS = {**SAFE, "deepseek-r1-1.5b": ("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", 2.2), "gemma3-4b": ("google/gemma-3-4b-it", 6.0)}
 
-BNB = BitsAndBytesConfig(
-  load_in_4bit=True, bnb_4bit_quant_type='nf4',
-  bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,
-)
+def gpu_ok():
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except: return False
 
-_cache = {}
+def gen(m, p, t):
+    if not gpu_ok():
+        return "ERROR: No GPU. Personal accounts need PRO for ZeroGPU. Use MCP client for 25min/day free, or run hf jobs for pay-as-you-go GPU. https://huggingface.co/mcp"
+    import spaces, torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    hid, vram = SPECS[m]
+    bnb = torch.classes if False else None  # dummy
+    bnb = None
+    try:
+        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
+    except: pass
+    tok = AutoTokenizer.from_pretrained(hid, trust_remote_code=True)
+    if tok.pad_token is None: tok.pad_token = tok.eos_token
+    kw = {"trust_remote_code": True, "low_cpu_mem_usage": True, "torch_dtype": torch.bfloat16}
+    if bnb: kw["quantization_config"] = bnb; kw["device_map"] = "auto"
+    mod = AutoModelForCausalLM.from_pretrained(hid, **kw)
+    msgs = [{"role":"user","content":p}]
+    txt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+    inp = tok(txt, return_tensors="pt", truncation=True, max_length=16384).to(next(mod.parameters()).device)
+    out = mod.generate(**inp, max_new_tokens=512, do_sample=t>0.01, temperature=max(t,1e-5), top_p=0.95, top_k=40, pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
+    return tok.decode(out[0], skip_special_tokens=True)
 
-def load_kv(hid):
-  if hid in _cache:
-    return _cache[hid]
-  tok = AutoTokenizer.from_pretrained(hid, trust_remote_code=True)
-  if tok.pad_token is None: tok.pad_token = tok.eos_token
-  mod = AutoModelForCausalLM.from_pretrained(
-    hid, quantization_config=BNB, device_map='auto',
-    torch_dtype=torch.bfloat16, trust_remote_code=True,
-    low_cpu_mem_usage=True, attn_implementation='sdpa',
-  )
-  _cache[hid] = (mod, tok)
-  return mod, tok
+with gr.Blocks(title="NEXUS LAB v3.4") as demo:
+    gr.Markdown("# NEXUS LAB v3.4")
+    gr.Markdown("**Status:** " + ("GPU Ready" if gpu_ok() else "CPU Only — No GPU detected"))
+    m = gr.Dropdown(choices=list(SPECS.keys()), value="llama3.2-1b", label="Model")
+    p = gr.Textbox(label="Prompt", lines=3)
+    t = gr.Slider(0.0, 2.0, 0.7, step=0.05, label="Temperature")
+    o = gr.Textbox(label="Output", lines=15)
+    def wrapped(m, p, t):
+        if not gpu_ok(): return "ERROR: No GPU. Enable ZeroGPU (PRO), use MCP for 25min/day free: https://huggingface.co/mcp, or run hf jobs."
+        return gen(m, p, t)
+    gr.Button("Generate", variant="primary").click(wrapped, [m, p, t], o)
+    gr.Markdown("---")
+    gr.Markdown("MCP: `/gradio_api/mcp/sse` — Use MCP client for free GPU via agents.md")
+    gr.Markdown("Jobs: `hf jobs run --flavor gpu-t4-small python train.py` — Pay-as-you-go, no PRO needed")
 
-@spaces.GPU(duration=lambda *a: 180, size='large')
-def gen(m, p, t, top_p):
-  if not p or not p.strip():
-    return 'Empty prompt.'
-  hid, vram, status = SPECS[m]
-  mod, tok = load_kv(hid)
-  msgs = [{'role':'user','content':p.strip()}]
-  txt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
-  inp = tok(txt, return_tensors='pt', truncation=True, max_length=16384).to(next(mod.parameters()).device)
-  out = mod.generate(
-    **inp, max_new_tokens=512, do_sample=t>0.01,
-    temperature=max(t,1e-5), top_p=top_p, top_k=40,
-    repetition_penalty=1.0, use_cache=True,
-    pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id,
-  )
-  r = tok.decode(out[0], skip_special_tokens=True)
-  return r
-
-with gr.Blocks(title='NEXUS LAB v3.3 - ZeroGPU') as demo:
-  gr.Markdown('## NEXUS LAB v3.3 - ZeroGPU Inference')
-  gr.Markdown('Enable ZeroGPU in Settings > Hardware. First load may take 1-2 minutes.')
-  m = gr.Dropdown(choices=list(SPECS.keys()), value='llama3.2-1b', label='Model')
-  p = gr.Textbox(label='Prompt', lines=3, value='Write a purple team strategy.')
-  t = gr.Slider(0.0, 2.0, 0.7, step=0.05, label='Temperature')
-  top_p = gr.Slider(0.1, 1.0, 0.95, step=0.05, label='Top-p')
-  o = gr.Textbox(label='Output', lines=15)
-  gr.Button('Generate', variant='primary').click(gen, [m, p, t, top_p], o)
-  gr.Markdown('MCP: /gradio_api/mcp/sse')
-demo.queue()
+demo.launch(mcp_server=True)