Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py — REFACTORED with a clean, custom Python loop using 'yield'
+# app.py — FIXED: Handles remote code trust and logical error on failure
 import gradio as gr
 import os
 import time
@@ -7,11 +7,11 @@ from transformers import LlavaForConditionalGeneration, AutoProcessor
 import torch
 
 # --- Config ---
-MODEL_NAME = "bczhou/TinyLLaVA-3.1B"  # or "llava-hf/llava-1.5-7b-hf"
+MODEL_NAME = "bczhou/TinyLLaVA-3.1B"
 CHECKPOINT_ROOT = "./checkpoints"
 os.makedirs(CHECKPOINT_ROOT, exist_ok=True)
 
-# --- Global state for the model (needed for the chat function) ---
+# --- Global state for the model ---
 current_stage = 0
 model = None
 processor = None
@@ -25,21 +25,28 @@ def load_model_for_stage(stage):
     """Loads the appropriate model and processor for a given stage."""
     global model, processor, current_stage
     
-    # Update the global stage so the chat function knows which model to use
     current_stage = stage
     
     ckpt_path = f"{CHECKPOINT_ROOT}/stage_{stage}"
+    # ✅ FIX 1: Added trust_remote_code=True to all .from_pretrained calls
     if os.path.exists(ckpt_path) and os.path.exists(os.path.join(ckpt_path, "adapter_model.safetensors")):
         print(f"✅ Loading checkpoint: Stage {stage}")
-        # Free up VRAM before loading the next model
         del model
         torch.cuda.empty_cache()
-        model = LlavaForConditionalGeneration.from_pretrained(ckpt_path, torch_dtype=torch.float16).to(device)
-        processor = AutoProcessor.from_pretrained(ckpt_path)
+        model = LlavaForConditionalGeneration.from_pretrained(
+            ckpt_path, 
+            torch_dtype=torch.float16, 
+            trust_remote_code=True
+        ).to(device)
+        processor = AutoProcessor.from_pretrained(ckpt_path, trust_remote_code=True)
     else:
         print(f"⚠️ No checkpoint for Stage {stage} — loading base model")
-        model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to(device)
-        processor = AutoProcessor.from_pretrained(MODEL_NAME)
+        model = LlavaForConditionalGeneration.from_pretrained(
+            MODEL_NAME, 
+            torch_dtype=torch.float16, 
+            trust_remote_code=True
+        ).to(device)
+        processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
 def chat_with_image(image, text, chat_history):
     """Handles the user's chat interaction."""
@@ -62,37 +69,34 @@ def chat_with_image(image, text, chat_history):
         chat_history.append({"role": "assistant", "content": f"⚠️ Error: {str(e)}"})
         return "", chat_history
 
-# --- The Custom Loop: Autonomous Training Pipeline ---
-# This single function runs the entire loop and 'yields' updates to the UI.
 def run_autonomous_training_and_update_ui():
     """
-    This is a generator function that runs the entire training pipeline.
-    It yields status messages that are displayed directly in the Gradio UI.
+    Generator function that runs the training pipeline and yields status updates.
     """
     yield "🚀 Initializing COCONUT-VLM Autonomous Trainer..."
+    
+    # ✅ FIX 2: Added a flag to track if training failed
+    training_failed = False
 
     for stage in [1, 2, 3]:
         ckpt_path = f"{CHECKPOINT_ROOT}/stage_{stage}"
         
-        # 1. Check if stage is already trained
         if os.path.exists(os.path.join(ckpt_path, "adapter_model.safetensors")):
             status_message = f"⏭️ Stage {stage} already trained — loading..."
             print(status_message)
             yield status_message
             load_model_for_stage(stage)
-            time.sleep(2) # Give user time to read the message
+            time.sleep(2)
             continue
 
-        # 2. Start training for the current stage
         status_message = f"▶️ AUTO-TRAINING STARTED: Stage {stage}"
         print(status_message)
         yield status_message
         
         try:
-            # This is the long-running training task
+            # IMPORTANT: Make sure train_vlm_stage also uses trust_remote_code=True
             train_vlm_stage(stage, MODEL_NAME, ckpt_path)
 
-            # 3. Handle successful training
             status_message = f"✅ Stage {stage} completed! Loading new model..."
             print(status_message)
             yield status_message
@@ -105,19 +109,19 @@ def run_autonomous_training_and_update_ui():
                 time.sleep(5)
 
         except Exception as e:
-            # 4. Handle training failure
-            status_message = f"❌ Stage {stage} failed: {str(e)}"
+            status_message = f"❌ Stage {stage} failed: {e}"
             print(status_message)
             yield status_message
-            break # Stop the entire pipeline if a stage fails
-
-    # 5. Final completion message
-    final_message = "🎉 COCONUT-VLM Training Complete — All 3 Stages Finished!"
-    print(final_message)
-    yield final_message
+            training_failed = True # Set the flag to True on failure
+            break # Stop the entire pipeline
 
+    # ��� FIX 2: Only show the completion message if the loop finished without failing
+    if not training_failed:
+        final_message = "🎉 COCONUT-VLM Training Complete — All 3 Stages Finished!"
+        print(final_message)
+        yield final_message
 
-# --- Gradio UI ---
+# --- Gradio UI (No changes needed here) ---
 with gr.Blocks(title="🥥 COCONUT-VLM Autonomous Trainer") as demo:
     gr.Markdown("# 🥥 COCONUT-VLM: Autonomous Vision-Language Trainer")
     gr.Markdown("Model is training itself in 3 stages automatically. **You can only chat.** Training is backend-only.")
@@ -129,7 +133,7 @@ with gr.Blocks(title="🥥 COCONUT-VLM Autonomous Trainer") as demo:
                 value="Waiting to start...",
                 interactive=False,
                 show_label=False,
-                lines=3 # Give it a bit more space
+                lines=3
             )
             gr.Markdown("💡 _Training runs automatically on page load. No buttons needed._")
 
@@ -139,14 +143,9 @@ with gr.Blocks(title="🥥 COCONUT-VLM Autonomous Trainer") as demo:
             msg = gr.Textbox(label="Ask a question about the image")
             clear = gr.Button("Clear Chat")
 
-    # Chat logic remains the same
     msg.submit(chat_with_image, [image_input, msg, chatbot], [msg, chatbot])
     clear.click(lambda: [], inputs=None, outputs=chatbot)
 
-    # --- THE MAGIC ---
-    # On page load, run our generator function. Gradio will automatically
-    # update the 'status' textbox every time the function 'yields' a new value.
-    # This is clean, efficient, and avoids all threading/polling headaches.
     demo.load(
         fn=run_autonomous_training_and_update_ui,
         inputs=None,