Merge pull request #369 from ebolam/Model_Plugins

Fix for --nobreakmodel forcing CPU

aiserver.py

@@ -1672,7 +1672,6 @@ class RestrictedUnpickler(pickle.Unpickler):
             )
 
     def load(self, *args, **kwargs):
-        logger.info("Using safe unpickle")
         self.original_persistent_load = getattr(
             self, "persistent_load", pickle.Unpickler.persistent_load
         )

modeling/inference_models/generic_hf_torch/class.py

@@ -250,7 +250,7 @@ class model_backend(HFTorchInferenceModel):
 
         
         if utils.koboldai_vars.hascuda:
-            if self.usegpu:
+            if self.usegpu or self.nobreakmodel:
                 # Use just VRAM
                 self.model = self.model.half().to(utils.koboldai_vars.gpu_device)
             elif self.breakmodel:

modeling/warpers.py

@@ -42,10 +42,17 @@ import utils
 import torch
 import numpy as np
 
-if utils.koboldai_vars.use_colab_tpu:
-    import jax
-    import jax.numpy as jnp
-    import tpu_mtj_backend
+try:
+    ignore = utils.koboldai_vars.use_colab_tpu
+    ok = True
+except:
+    ok = False
+
+if ok:
+    if utils.koboldai_vars.use_colab_tpu:
+        import jax
+        import jax.numpy as jnp
+        import tpu_mtj_backend
 
 
 def update_settings():