Fix for --nobreakmodel forcing CPU

Put importing of colab packages into a if function so it doesn't error out

aiserver.py

@@ -1682,7 +1682,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,6 +42,13 @@ import utils
 import torch
 import numpy as np
 
+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