Upload app.py

app.py

@@ -326,12 +326,18 @@ def load_model():
             torch_dtype=torch.bfloat16,
             trust_remote_code=False
         ).to(device).eval()
-    # Compile model for faster inference (reduce-overhead is safest for generative loops)
-    try:
-        model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
-        print("Model compiled with torch.compile.")
-    except Exception as e:
-        print(f"torch.compile skipped: {e}")
+    # Compile model for faster inference — ONLY for standard causal models
+    if not IS_DIFFUSION:
+        try:
+            model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
+            print("Model compiled with torch.compile.")
+        except Exception as e:
+            print(f"torch.compile skipped: {e}")
+    else:
+        print("Diffusion model loaded without torch.compile (custom FX code incompatible with Dynamo).")
+
+
+
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION