Fix ZeroGPU adapter loading + add CPU fallback

app.py

@@ -16,14 +16,26 @@ LANG_META = {
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
 
-model = PeftModel.from_pretrained(base_model, LANG_META["Python"]["repo"], adapter_name="Python")
+model = PeftModel.from_pretrained(
+    base_model, LANG_META["Python"]["repo"], adapter_name="Python", device_map="cpu"
+)
 for lang, meta in LANG_META.items():
     if lang != "Python":
-        model.load_adapter(meta["repo"], adapter_name=lang)
+        model.load_adapter(meta["repo"], adapter_name=lang, device_map="cpu")
 model.eval()
 
 _on_gpu = False
 
+
+def _generate(inputs, use_cuda):
+    device = "cuda" if use_cuda else "cpu"
+    model.to(device)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    return model.generate(
+        **inputs, max_new_tokens=512, temperature=0.2, do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
 EXAMPLES = [
     ["Merge k sorted linked lists into one sorted list.", "Divide and conquer / heap", "Python"],
     ["Given a set of non-overlapping intervals, insert a new interval and merge as needed.", "Sorting / interval merge", "JavaScript"],
@@ -116,9 +128,6 @@ THEME = gr.themes.Base(
 @spaces.GPU(duration=60)
 def solve(problem, algorithm_tag, language):
     global _on_gpu
-    if not _on_gpu:
-        model.to("cuda")
-        _on_gpu = True
     if not problem.strip():
         return "", "Enter a problem statement first."
 
@@ -135,14 +144,21 @@ def solve(problem, algorithm_tag, language):
 
     messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_msg}]
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt")
 
-    outputs = model.generate(
-        **inputs, max_new_tokens=512, temperature=0.2, do_sample=True,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-    code = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-    return code, f"{LANG_META[language]['icon']} generated with the {language} QDoRA adapter"
+    try:
+        outputs = _generate(inputs, use_cuda=True)
+        _on_gpu = True
+        engine_note = "GPU"
+    except RuntimeError as e:
+        if "CUDA" not in str(e) and "cuda" not in str(e):
+            raise
+        outputs = _generate(inputs, use_cuda=False)
+        engine_note = "CPU fallback"
+
+    input_len = inputs["input_ids"].shape[1]
+    code = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
+    return code, f"{LANG_META[language]['icon']} generated with the {language} QDoRA adapter · {engine_note}"
 
 
 def on_lang_change(language):