Update flux_train_ui.py

flux_train_ui.py

@@ -3,6 +3,11 @@ from huggingface_hub import whoami
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 import sys
 import spaces
+import gradio as gr
+from fastapi import FastAPI, Request
+import uvicorn
+from fastapi import Request
+
 # Add the current working directory to the Python path
 sys.path.insert(0, os.getcwd())
 
@@ -21,6 +26,26 @@ sys.path.insert(0, "ai-toolkit")
 from toolkit.job import get_job
 
 MAX_IMAGES = 150
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 
 def load_captioning(uploaded_files, concept_sentence):
     uploaded_images = [file for file in uploaded_files if not file.endswith('.txt')]
@@ -427,5 +452,34 @@ with gr.Blocks(theme=theme, css=css) as demo:
 
     do_captioning.click(fn=run_captioning, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)
 
+
+
+
+
+
+
+def train_model_ui(data: str):
+    print(f"🚀 Triggered with input: {data}")
+    return f"✅ Received: {data}"
+    
+def add_trigger_endpoint(app):
+    @app.post("/trigger")
+    async def trigger(request: Request):
+        try:
+            body = await request.json()
+            input_data = body.get("input", "")
+            print(f"🔁 API Trigger: {input_data}")
+            result = train_model_ui(input_data)
+            return {"result": result}
+        except Exception as e:
+            return {"error": str(e)}
+
+# ⬇️ DO NOT remove this since you're keeping the manual launch
 if __name__ == "__main__":
-    demo.launch(share=True, show_error=True)
+    # Launch Gradio, then hook in the FastAPI route
+    demo.launch(
+        share=True,
+        show_error=True,
+        after_start=add_trigger_endpoint  # ✅ Hook the FastAPI logic here
+    )
+