Upload mod.py

mod.py

@@ -13,7 +13,16 @@ from env import num_cns, model_trigger, HF_TOKEN, CIVITAI_API_KEY, DOWNLOAD_LORA
 from modutils import download_things
 
 
-subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+IS_ZERO = True if os.getenv("SPACES_ZERO_GPU", None) else False
+if torch.cuda.is_available():
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cuda.matmul.allow_tf32 = True
+    if IS_ZERO:
+        subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+        torch.set_float32_matmul_precision("high") # https://pytorch.org/blog/accelerating-generative-ai-3/
+
+
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 #subprocess.run('pip cache purge', shell=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"