From d426961266e645bd515d44c34ead62a85b60406b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 1 Jan 2025 16:34:18 -0800
Subject: [PATCH] Rewind

---
 unsloth/models/_utils.py |  4 +--
 unsloth/models/loader.py | 60 +++++++++++++++++++---------------------
 2 files changed, 31 insertions(+), 33 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 3cb6ffb8f3..386d71354d 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1203,8 +1203,6 @@ def unsloth_compile_transformers(
         return
     pass
 
-    if disable: return
-
     model_types = get_transformers_model_type(
         model_name        = model_name,
         token             = token,
@@ -1212,6 +1210,8 @@ def unsloth_compile_transformers(
         trust_remote_code = trust_remote_code,
     )
 
+    if disable: return
+
     for model_type in model_types:
         _unsloth_compile_transformers(
             model_type,
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 20c0177d7e..2ec7745154 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -454,37 +454,35 @@ class FastVisionModel(FastBaseVisionModel):
 
         if not was_disabled: enable_progress_bars()
 
-        if os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") == "0":
-            with contextlib.redirect_stdout(open(os.devnull, "w")):
-                patch_loss_functions(torch_compile = False)
-                model_types = unsloth_compile_transformers(
-                    model_name              = model_name,
-                    sdpa_dynamic_mask       = True,
-                    sdpa_bool_masks         = True,
-                    sdpa_gqa_replace        = True,
-                    sdpa_dynamic_compile    = True,
-                    compile_attention       = True,
-                    disable_causal_masks    = True,
-                    compile_torch_modules   = True,
-                    compile_custom_modules  = True,
-                    compile_function_calls  = True,
-                    fuse_lm_head            = True,
-                    gradient_checkpointing  = True,
-                    manual_replacements     = True,
-                    fast_lora_forwards      = True,
-                    fast_residual_stream    = False,
-                    accurate_accumulation   = True,
-                    epilogue_fusion         = True,
-                    max_autotune            = False,
-                    shape_padding           = True,
-                    cudagraphs              = False,
-                    debug                   = False,
-                    fullgraph               = fullgraph,
-                    import_from_cache       = False,
-                    disable                 = False,
-                    return_logits           = return_logits,
-                )
-            pass
+        with contextlib.redirect_stdout(open(os.devnull, "w")):
+            patch_loss_functions(torch_compile = False)
+            model_types = unsloth_compile_transformers(
+                model_name              = model_name,
+                sdpa_dynamic_mask       = True,
+                sdpa_bool_masks         = True,
+                sdpa_gqa_replace        = True,
+                sdpa_dynamic_compile    = True,
+                compile_attention       = True,
+                disable_causal_masks    = True,
+                compile_torch_modules   = True,
+                compile_custom_modules  = True,
+                compile_function_calls  = True,
+                fuse_lm_head            = True,
+                gradient_checkpointing  = True,
+                manual_replacements     = True,
+                fast_lora_forwards      = True,
+                fast_residual_stream    = False,
+                accurate_accumulation   = True,
+                epilogue_fusion         = True,
+                max_autotune            = False,
+                shape_padding           = True,
+                cudagraphs              = False,
+                debug                   = False,
+                fullgraph               = fullgraph,
+                import_from_cache       = False,
+                disable                 = False,
+                return_logits           = return_logits,
+            )
         pass
 
         # Check if this is local model since the tokenizer gets overwritten