Update app.py

app.py

@@ -115,9 +115,39 @@ if __name__ == "__main__":
     import gradio as gr
     from datetime import datetime
 
+    def start_training(instance_token, class_token, zip_file, output_dir, max_steps, lr, hf_token, seed, precision):
+        return train_model(
+            instance_token=instance_token,
+            class_token=class_token,
+            zip_path=zip_file.name,
+            output_dir=output_dir,
+            max_train_steps=max_steps,
+            learning_rate=lr,
+            hf_token=hf_token,
+            seed=seed,
+            precision=precision
+        )
+
     def create_ui():
         with gr.Blocks() as demo:
-            gr.Textbox(label="App Initialized")
+            with gr.Tab("Train Model"):
+                instance_token = gr.Textbox(label="Instance Token")
+                class_token = gr.Textbox(label="Class Token")
+                zip_file = gr.File(label="Training ZIP File")
+                output_dir = gr.Textbox(label="Output Directory", value="trained_model")
+                max_steps = gr.Number(label="Max Training Steps", value=1200)
+                lr = gr.Number(label="Learning Rate", value=5e-6)
+                seed = gr.Number(label="Random Seed", value=42)
+                precision = gr.Dropdown(label="Precision Mode", choices=["fp16", "fp32"], value="fp16")
+                hf_token_train = gr.Textbox(label="Hugging Face Token", type="password")
+                train_btn = gr.Button("Start Training")
+                train_output = gr.Textbox(label="Training Output", lines=8)
+
+                train_btn.click(
+                    fn=start_training,
+                    inputs=[instance_token, class_token, zip_file, output_dir, max_steps, lr, hf_token_train, seed, precision],
+                    outputs=train_output
+                )
         return demo
 
     demo = create_ui()