Update app.py

app.py

@@ -14,11 +14,19 @@ MODEL_BASE_THINK = os.path.join(MODEL_BASE_SINGLE, 'stage35')
 
 # model_single = llava.load(MODEL_BASE_SINGLE, model_base=None, devices=[0])
 model_single = llava.load(MODEL_BASE_SINGLE, model_base=None)
-model_single = model_single.to("cuda")
-model_single_copy = copy.deepcopy(model_single)
+model_single = model_single.to("cuda") 
+model_single_copy = copy.deepcopy(model_single) # keep a copy of the original model for non-thinking mode
 
 generation_config_single = model_single.default_generation_config
 
+non_lora_trainables = torch.load(
+                os.path.join(MODEL_BASE_THINK, "non_lora_trainables.bin"),
+                map_location="cpu",
+            )
+non_lora_trainables = {
+        (k[6:] if k.startswith("model.") else k): v for k, v in non_lora_trainables.items()
+    }
+model.load_state_dict(non_lora_trainables, strict=False)
 model_think = PeftModel.from_pretrained(
         model_single,
         MODEL_BASE_THINK,