Retry AoTI and attention backend; show full errors

app.py

@@ -216,6 +216,11 @@ def set_attention_backend_safe(transformer, backend: str) -> str:
                         break
                 if attention_backend_error is None and last_exc is not None:
                     attention_backend_error = str(last_exc)
+                if attention_backend_error:
+                    print(
+                        f"Requested attention backend {backend!r} failed; using {name!r} instead. "
+                        f"Reason: {attention_backend_error}"
+                    )
             return name
         except Exception as exc:  # noqa: BLE001
             last_exc = exc
@@ -538,26 +543,27 @@ def warmup_model(pipeline: ZImagePipeline, resolutions: List[str]) -> None:
 
 
 def init_app() -> None:
-    global aoti_loaded, aoti_error
+    global aoti_loaded, aoti_error, applied_attention_backend
     try:
         ensure_models_loaded()
         if ENABLE_COMPILE and pipe is not None:
             ensure_on_gpu()
         if ENABLE_AOTI and not aoti_loaded and pipe is not None and getattr(pipe, "transformer", None) is not None:
-            if not module_available("kernels"):
+            try:
+                pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
+                spaces.aoti_blocks_load(pipe.transformer.layers, AOTI_REPO, variant=AOTI_VARIANT)
+                aoti_loaded = True
+                aoti_error = None
+                print(f"AoTI loaded: {AOTI_REPO} (variant={AOTI_VARIANT})")
+            except Exception as exc:  # noqa: BLE001
                 aoti_loaded = False
-                aoti_error = "kernels module not available"
-                print("AoTI unavailable (kernels module not available).")
-            else:
-                try:
-                    pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-                    spaces.aoti_blocks_load(pipe.transformer.layers, AOTI_REPO, variant=AOTI_VARIANT)
-                    aoti_loaded = True
-                    aoti_error = None
-                    print(f"AoTI loaded: {AOTI_REPO} (variant={AOTI_VARIANT})")
-                except Exception as exc:  # noqa: BLE001
-                    aoti_error = str(exc)
-                    print(f"AoTI load failed (continuing without AoTI): {exc}")
+                aoti_error = str(exc)
+                print(f"AoTI load failed (continuing without AoTI): {exc}")
+            try:
+                applied_attention_backend = set_attention_backend_safe(pipe.transformer, ATTENTION_BACKEND)
+                print(f"Attention backend (post-AoTI): {applied_attention_backend}")
+            except Exception as exc:  # noqa: BLE001
+                print(f"Attention backend update failed (continuing): {exc}")
         if ENABLE_WARMUP and pipe is not None:
             ensure_on_gpu()
             try:
@@ -670,6 +676,14 @@ Attention: `{attention_status}` | AoTI: `{aoti_status}` | torch.compile: `{compi
 
 </div>"""
     )
+    if attention_backend_error or aoti_error:
+        with gr.Accordion("Runtime details (debug)", open=False):
+            details: List[str] = []
+            if attention_backend_error:
+                details.append(f"**Attention backend error**\n```\n{attention_backend_error}\n```")
+            if aoti_error:
+                details.append(f"**AoTI error**\n```\n{aoti_error}\n```")
+            gr.Markdown("\n\n".join(details))
 
     with gr.Row():
         with gr.Column(scale=1):