Update flux_train_ui.py

flux_train_ui.py

@@ -466,9 +466,18 @@ def train_model_ui(data: str):
 
 # ⬇️ DO NOT remove this since you're keeping the manual launch
 if __name__ == "__main__":
-    app = demo.launch(share=True, show_error=True, prevent_thread_lock=True)
+    # Unpack launch return value if it's a tuple
+    result = demo.launch(share=True, show_error=True, prevent_thread_lock=True)
 
-    @app.post("/trigger")
+    if isinstance(result, tuple):
+        server = result[0]
+    else:
+        server = result
+
+    # Try to get the FastAPI app object
+    fastapi_app = getattr(server, "app", None) or getattr(server, "server", None).app
+
+    @fastapi_app.post("/trigger")
     async def trigger(request: Request):
         try:
             body = await request.json()
@@ -478,7 +487,7 @@ if __name__ == "__main__":
         except Exception as e:
             return {"error": str(e)}
 
-    app.block()  # keep the server running
+    server.block_thread()