Update app.py

app.py

@@ -4,14 +4,12 @@ import os
 import time
 from train_vlm import train_vlm_stage
 from transformers import AutoImageProcessor, AutoTokenizer
-from custom_vlm import CustomScratchVLM # Import our custom model
+from custom_vlm import CustomScratchVLM
 import torch
 
-# --- Config ---
 CHECKPOINT_ROOT = "./checkpoints"
 os.makedirs(CHECKPOINT_ROOT, exist_ok=True)
 
-# --- Global state for the model ---
 current_stage = 0
 model = None
 image_processor = None
@@ -19,17 +17,13 @@ tokenizer = None
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print(f"🖥️ Running on device: {device}")
-if device == "cuda":
-    print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
+if device == "cuda": print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
 
 def load_model_for_stage(stage):
-    """Loads the appropriate custom model and processors for a given stage."""
     global model, image_processor, tokenizer, current_stage
-    
     current_stage = stage
     ckpt_path = f"{CHECKPOINT_ROOT}/stage_{stage}"
     
-    # Check for a saved model config, which indicates a trained checkpoint
     if os.path.exists(os.path.join(ckpt_path, "config.json")):
         print(f"✅ Loading FROM-SCRATCH checkpoint: Stage {stage}")
         if model is not None: del model
@@ -39,103 +33,79 @@ def load_model_for_stage(stage):
         image_processor = AutoImageProcessor.from_pretrained(ckpt_path)
         tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
     else:
-        # Before any training, there is no model to load.
         print(f"⚠️ No checkpoint for Stage {stage} — model is not loaded.")
-        model = None
-        image_processor = None
-        tokenizer = None
+        model, image_processor, tokenizer = None, None, None
 
 def chat_with_image(image, text, chat_history):
-    """Handles the user's chat interaction with the custom VLM."""
-    if model is None or image_processor is None or tokenizer is None:
-        return "", chat_history + [{"role": "assistant", "content": "Model is not loaded or is currently training from scratch. Please wait."}]
+    if not all([model, image_processor, tokenizer]):
+        return "", chat_history + [{"role": "assistant", "content": "Model is not loaded or is currently training. Please wait."}]
     
     if image is None:
         return "", chat_history + [{"role": "user", "content": text}, {"role": "assistant", "content": "Please upload an image."}]
 
     try:
-        # Prepare inputs for our custom model
         pixel_values = image_processor(image, return_tensors="pt").pixel_values.to(device)
         
-        # Format prompt with the special image token
-        prompt = f"USER: {tokenizer.additional_special_tokens[0]}\nQ: {text}\nASSISTANT:"
-        prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        # For inference, we do not include the <IMAGE> token in the text prompt
+        prompt = f"USER: \nQuestion: {text}\nASSISTANT:"
+        inputs = tokenizer(prompt, return_tensors="pt")
+        input_ids = inputs.input_ids.to(device)
+        attention_mask = inputs.attention_mask.to(device)
 
-        # Use our custom generate method
         output_ids = model.generate(
             pixel_values=pixel_values, 
-            prompt_ids=prompt_ids,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=256,
             do_sample=True,
-            temperature=0.7
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
         )
         
-        # Decode the generated part
-        response = tokenizer.decode(output_ids[0][prompt_ids.shape[1]:], skip_special_tokens=True)
+        # Decode only the newly generated tokens
+        response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
 
         chat_history.append({"role": "user", "content": text})
-        chat_history.append({"role": "assistant", "content": response})
+        chat_history.append({"role": "assistant", "content": response or "[No response generated]"})
         return "", chat_history
     except Exception as e:
-        return "", chat_history + [{"role": "user", "content": text}, {"role": "assistant", "content": f"⚠️ Error: {str(e)}"}]
-
+        return "", chat_history + [{"role": "user", "content": text}, {"role": "assistant", "content": f"⚠️ Error: {e}"}]
 
 def run_autonomous_training_and_update_ui():
-    """Generator function that runs the from-scratch training pipeline."""
-    yield "🚀 Initializing COCONUT-VLM Autonomous Trainer (FROM SCRATCH)..."
-    
-    training_failed = False
+    yield "🚀 Initializing From-Scratch Trainer..."
     for stage in [1, 2, 3]:
         ckpt_path = f"{CHECKPOINT_ROOT}/stage_{stage}"
-        
         if os.path.exists(os.path.join(ckpt_path, "config.json")):
             status_message = f"⏭️ Stage {stage} already trained — loading..."
             yield status_message
             load_model_for_stage(stage)
             continue
 
-        status_message = f"▶️ AUTO-TRAINING FROM SCRATCH STARTED: Stage {stage}"
+        status_message = f"▶️ AUTO-TRAINING FROM SCRATCH: Stage {stage}"
         yield status_message
-        
         try:
-            # Call the training function (no longer needs model_name)
             train_vlm_stage(stage, ckpt_path)
             status_message = f"✅ Stage {stage} completed! Loading new model..."
             yield status_message
             load_model_for_stage(stage)
         except Exception as e:
             status_message = f"❌ Stage {stage} failed: {e}"
-            yield status_message
-            training_failed = True
-            break
+            yield status_message; raise e # Stop execution on failure
+    yield "🎉 COCONUT-VLM Training Complete — All 3 Stages Finished!"
 
-    if not training_failed:
-        yield "🎉 COCONUT-VLM Training Complete — All 3 Stages Finished!"
-    else:
-        yield "🛑 Training stopped due to an error."
-
-# --- Gradio UI (No Changes) ---
-with gr.Blocks(title="🥥 COCONUT-VLM Autonomous Trainer") as demo:
-    gr.Markdown("# 🥥 COCONUT-VLM: Autonomous Vision-Language Trainer (From Scratch)")
-    gr.Markdown("Model is training itself **from random initialization**. This will take a very long time and requires significant compute. You can interact with the latest trained model.")
-    # ... rest of UI is the same ...
+with gr.Blocks(title="🥥 COCONUT-VLM From Scratch") as demo:
+    gr.Markdown("# 🥥 COCONUT-VLM: Autonomous Trainer (From Scratch)")
+    gr.Markdown("Model is training itself **from random initialization**. You can interact with the latest trained model.")
     with gr.Row():
         with gr.Column(scale=1):
-            status = gr.Textbox(
-                label="Training Status", value="Waiting to start...", interactive=False,
-                lines=10, max_lines=20
-            )
-            gr.Markdown("💡 _Training runs automatically on page load._")
-
+            status = gr.Textbox(label="Training Status", value="Waiting to start...", interactive=False, lines=10)
         with gr.Column(scale=2):
             image_input = gr.Image(type="pil", label="Upload Image")
             chatbot = gr.Chatbot(label="Chat with the VLM", height=400)
-            msg = gr.Textbox(label="Ask a question about the image")
+            msg = gr.Textbox(label="Ask a question")
             clear = gr.Button("Clear Chat")
-
     msg.submit(chat_with_image, [image_input, msg, chatbot], [msg, chatbot])
     clear.click(lambda: (None, None, []), None, [image_input, msg, chatbot])
-
     demo.load(fn=run_autonomous_training_and_update_ui, inputs=None, outputs=status)
 
 demo.queue().launch(debug=True)