Update flux_train_ui.py

flux_train_ui.py

@@ -466,18 +466,9 @@ def train_model_ui(data: str):
 
 # ⬇️ DO NOT remove this since you're keeping the manual launch
 if __name__ == "__main__":
-    # Unpack launch return value if it's a tuple
-    result = demo.launch(share=True, show_error=True, prevent_thread_lock=True)
+    app = demo.launch(share=True, show_error=True, prevent_thread_lock=True)
 
-    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")
+    @app.post("/trigger")
     async def trigger(request: Request):
         try:
             body = await request.json()
@@ -487,7 +478,7 @@ if __name__ == "__main__":
         except Exception as e:
             return {"error": str(e)}
 
-    server.block_thread()
+    app.block()  # keep the server running