Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py — Fully autonomous 3-stage VLM trainer. UI is chat-only.
+# app.py — FIXED: Gradio 4.x compatible, no deprecation warnings, auto-trains stages
 import gradio as gr
 import threading
 import os
@@ -19,10 +19,14 @@ processor = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 training_status = "🚀 Initializing COCONUT-VLM Autonomous Trainer..."
 
+print(f"🖥️ Running on device: {device}")
+if device == "cuda":
+    print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
+
 def load_model_for_stage(stage):
     global model, processor
     ckpt_path = f"{CHECKPOINT_ROOT}/stage_{stage}"
-    if os.path.exists(ckpt_path):
+    if os.path.exists(ckpt_path) and os.path.exists(os.path.join(ckpt_path, "adapter_model.safetensors")):
         print(f"✅ Loading checkpoint: Stage {stage}")
         model = LlavaForConditionalGeneration.from_pretrained(ckpt_path, torch_dtype=torch.float16).to(device)
         processor = AutoProcessor.from_pretrained(ckpt_path)
@@ -36,6 +40,7 @@ def chat_with_image(image, text, chat_history):
         load_model_for_stage(current_stage)
 
     try:
+        # Format input for model
         conversation = [
             {"role": "user", "content": f"<image>\n{text}"},
         ]
@@ -45,10 +50,13 @@ def chat_with_image(image, text, chat_history):
         output = model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
         response = processor.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
 
-        chat_history.append((text, response))
+        # Append as OpenAI-style messages (fixes deprecation warning)
+        chat_history.append({"role": "user", "content": text})
+        chat_history.append({"role": "assistant", "content": response})
         return "", chat_history
     except Exception as e:
-        chat_history.append((text, f"⚠️ Error: {str(e)}"))
+        chat_history.append({"role": "user", "content": text})
+        chat_history.append({"role": "assistant", "content": f"⚠️ Error: {str(e)}"})
         return "", chat_history
 
 # --- Autonomous Training Pipeline ---
@@ -60,18 +68,22 @@ def auto_train_pipeline():
         training_status = f"▶️ AUTO-TRAINING STARTED: Stage {stage}"
         print(training_status)
 
+        ckpt_path = f"{CHECKPOINT_ROOT}/stage_{stage}"
+        # Skip if already trained
+        if os.path.exists(os.path.join(ckpt_path, "adapter_model.safetensors")):
+            training_status = f"⏭️ Stage {stage} already trained — loading..."
+            print(training_status)
+            load_model_for_stage(stage)
+            time.sleep(3)
+            continue
+
         try:
-            # Train stage
-            train_vlm_stage(stage, MODEL_NAME, f"{CHECKPOINT_ROOT}/stage_{stage}")
+            train_vlm_stage(stage, MODEL_NAME, ckpt_path)
 
-            # Update status
             training_status = f"✅ Stage {stage} completed! Loading model..."
             print(training_status)
-
-            # Load newly trained model
             load_model_for_stage(stage)
 
-            # Brief pause before next stage
             if stage < 3:
                 training_status = f"⏳ Advancing to Stage {stage + 1} in 5 seconds..."
                 print(training_status)
@@ -80,7 +92,7 @@ def auto_train_pipeline():
         except Exception as e:
             training_status = f"❌ Stage {stage} failed: {str(e)}"
             print(training_status)
-            break  # Stop pipeline on failure
+            break
 
     training_status = "🎉 COCONUT-VLM Training Complete — All 3 Stages Finished!"
     print(training_status)
@@ -98,20 +110,23 @@ with gr.Blocks(title="🥥 COCONUT-VLM Autonomous Trainer") as demo:
     with gr.Row():
         with gr.Column(scale=1):
             status = gr.Textbox(label="Training Status", value="Initializing...", interactive=False)
-            gr.Markdown("💡 _Training runs automatically in background. No buttons. No switching._")
+            gr.Markdown("💡 _Training runs automatically. No buttons. No switching._")
 
         with gr.Column(scale=2):
             image_input = gr.Image(type="pil", label="Upload Image")
-            chatbot = gr.Chatbot(height=400)
+
+            # ✅ FIXED: Set type="messages" to avoid deprecation warning
+            chatbot = gr.Chatbot(height=400, type="messages")
+
             msg = gr.Textbox(label="Ask a question about the image")
             clear = gr.Button("Clear Chat")
 
+    # Chat logic
     msg.submit(chat_with_image, [image_input, msg, chatbot], [msg, chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
-    # Initialize autonomous training on launch
+    # ✅ FIXED: Use Gradio 4.x compatible .load() with every=
     demo.load(initialize_autonomous_trainer, inputs=None, outputs=None)
-    # Poll training status every 3 seconds
-    demo.load(lambda: training_status, every=3, outputs=status)
+    demo.load(lambda: training_status, inputs=None, outputs=status, every=3)  # ← Now compatible
 
 demo.queue(max_size=20).launch()