Update flux_train_ui.py

flux_train_ui.py

@@ -465,29 +465,22 @@ 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)
-
-    if isinstance(result, tuple):
-        server = result[0]
-    else:
-        server = result
+# Main FastAPI app
+app = FastAPI()
 
-    # Try to get the FastAPI app object
-    fastapi_app = getattr(server, "app", None) or getattr(server, "server", None).app
+# Mount the Gradio app at root ("/")
+app = gr.mount_gradio_app(app, demo, path="/")
 
-    @fastapi_app.post("/trigger")
-    async def trigger(request: Request):
-        try:
-            body = await request.json()
-            input_data = body.get("input", "")
-            result = train_model_ui(input_data)
-            return {"result": result}
-        except Exception as e:
-            return {"error": str(e)}
-
-    server.block_thread()
+# Custom POST endpoint
+@app.post("/trigger")
+async def trigger(request: Request):
+    body = await request.json()
+    input_data = body.get("input", "")
+    result = train_model_ui(input_data)
+    return {"result": result}
 
+# Run everything
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)