Update app.py

app.py

@@ -57,35 +57,36 @@ pipe.to("cuda")
 # ---------------------------------------------------------------------------
 # 🔥 AOT LOADING LOGIC 🔥
 # ---------------------------------------------------------------------------
-AOT_FILENAME = "helios_distilled_transformer.pt2"
-AOT_PATH = os.path.join(_APP_DIR, AOT_FILENAME)
-
-def load_aot_model(path, original_module):
-    """
-    Loads a raw AOTI package (.pt2) and patches the original module.
-    """
-    print(f"[AOT] Loading AOTI package from {path}...")
-    
-    compiled_model = torch._inductor.aoti_load_package(path)
-    
-    original_module.forward = compiled_model
-    
-    original_module.to("meta") 
-    
-    print("[AOT] Model patched successfully!")
-
-if os.path.exists(AOT_PATH):
-    try:
-        load_aot_model(AOT_PATH, pipe.transformer)
-        print(f"[AOT] ✅ Loaded compiled graph: {e}")
-    except Exception as e:
-        print(f"[AOT] ❌ Failed to load compiled graph: {e}")
-        # Restore device if failed
-        pipe.to("cuda")
-        pipe.transformer.set_attention_backend("_flash_3_hub")
-else:
-    print(f"[AOT] ⚠️ No compiled graph found at {AOT_PATH}.")
-    pipe.transformer.set_attention_backend("_flash_3_hub")
+# AOT_FILENAME = "helios_distilled_transformer.pt2"
+# AOT_PATH = os.path.join(_APP_DIR, AOT_FILENAME)
+
+#def load_aot_model(path, original_module):
+#    """
+#    Loads a raw AOTI package (.pt2) and patches the original module.
+#    """
+#    print(f"[AOT] Loading AOTI package from {path}...")
+#    
+#    compiled_model = torch._inductor.aoti_load_package(path)
+#    
+#    original_module.forward = compiled_model
+#    
+#    original_module.to("meta") 
+#    
+#    print("[AOT] Model patched successfully!")
+
+#if os.path.exists(AOT_PATH):
+#    try:
+#        load_aot_model(AOT_PATH, pipe.transformer)
+#        print(f"[AOT] ✅ Loaded compiled graph: {e}")
+#    except Exception as e:
+#        print(f"[AOT] ❌ Failed to load compiled graph: {e}")
+#        # Restore device if failed
+#        pipe.to("cuda")
+#        pipe.transformer.set_attention_backend("_flash_3_hub")
+#else:
+#    print(f"[AOT] ⚠️ No compiled graph found at {AOT_PATH}.")
+
+pipe.transformer.set_attention_backend("_flash_3_hub")
 
 # ---------------------------------------------------------------------------
 # Generation