diff --git a/demo.ipynb b/demo.ipynb index 019bcb2..cb8a74d 100644 --- a/demo.ipynb +++ b/demo.ipynb @@ -567,50 +567,62 @@ "metadata": {}, "outputs": [], "source": [ + "import torch\n", + "import gc\n", + "\n", "# Ensure the model is reset to the clean base model state (LoRA hooks disabled)\n", "model.reset()\n", + "torch.cuda.empty_cache()\n", + "gc.collect()\n", "\n", "base_predictions = []\n", "print(\"Running Base Model Inference (with visual tokens) for all 6 examples...\")\n", - "for idx, example in enumerate(examples):\n", - " print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n", - " \n", - " base_messages = [\n", - " [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": [\n", - " {\"type\": \"video\", \"path\": example[\"video_path\"]},\n", - " {\"type\": \"text\", \"text\": example[\"prompt\"]}\n", - " ]\n", - " }\n", + "\n", + "with torch.inference_mode():\n", + " for idx, example in enumerate(examples):\n", + " print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n", + " \n", + " base_messages = [\n", + " [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video\", \"path\": example[\"video_path\"]},\n", + " {\"type\": \"text\", \"text\": example[\"prompt\"]}\n", + " ]\n", + " }\n", + " ]\n", " ]\n", - " ]\n", "\n", - " base_inputs = prepare_smolvlm_inputs(\n", - " processor,\n", - " base_messages,\n", - " device,\n", - " video_fps=train_args.video_fps,\n", - " max_frames=train_args.max_frames,\n", - " video_size_longest_edge=train_args.video_size_longest_edge,\n", - " video_load_backend=train_args.video_load_backend\n", - " )\n", + " base_inputs = prepare_smolvlm_inputs(\n", + " processor,\n", + " base_messages,\n", + " device,\n", + " video_fps=train_args.video_fps,\n", + " max_frames=train_args.max_frames,\n", + " video_size_longest_edge=train_args.video_size_longest_edge,\n", + " video_load_backend=train_args.video_load_backend\n", + " )\n", "\n", - " base_generated_ids = raw_model.generate(\n", - " **base_inputs,\n", - " max_new_tokens=train_args.generation_max_new_tokens,\n", - " do_sample=False,\n", - " pad_token_id=tokenizer.pad_token_id,\n", - " eos_token_id=tokenizer.eos_token_id\n", - " )\n", + " base_generated_ids = raw_model.generate(\n", + " **base_inputs,\n", + " max_new_tokens=train_args.generation_max_new_tokens,\n", + " do_sample=False,\n", + " pad_token_id=tokenizer.pad_token_id,\n", + " eos_token_id=tokenizer.eos_token_id\n", + " )\n", "\n", - " base_pred = tokenizer.decode(\n", - " base_generated_ids[0][base_inputs[\"input_ids\"].shape[1]:],\n", - " skip_special_tokens=True\n", - " ).strip()\n", - " \n", - " base_predictions.append(base_pred)\n", + " base_pred = tokenizer.decode(\n", + " base_generated_ids[0][base_inputs[\"input_ids\"].shape[1]:],\n", + " skip_special_tokens=True\n", + " ).strip()\n", + " \n", + " base_predictions.append(base_pred)\n", + " \n", + " # Clean up temporary GPU tensors for this step\n", + " del base_inputs, base_generated_ids\n", + " torch.cuda.empty_cache()\n", + " gc.collect()\n", "\n", "print(\"\\nBase Model inferences completed successfully.\")" ] @@ -634,105 +646,120 @@ "metadata": {}, "outputs": [], "source": [ + "import torch\n", + "import gc\n", + "\n", "video2lora_predictions = []\n", "print(\"Running Video2LoRA Internalization & Inference (zero visual tokens) for all 6 examples...\")\n", - "for idx, example in enumerate(examples):\n", - " print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n", - " \n", - " # 1. Reset model to ensure clean base features are extracted\n", - " model.reset()\n", - " \n", - " # 2. Extract context activations\n", - " internalize_messages = [\n", - " [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": [\n", - " {\"type\": \"video\", \"path\": example[\"video_path\"]},\n", - " {\"type\": \"text\", \"text\": train_args.internalization_prompt}\n", - " ]\n", - " }\n", + "\n", + "with torch.inference_mode():\n", + " for idx, example in enumerate(examples):\n", + " print(f\"Processing Example {idx+1}/{len(examples)}: {example['id']}...\")\n", + " \n", + " # 1. Reset model to ensure clean base features are extracted\n", + " model.reset()\n", + " torch.cuda.empty_cache()\n", + " gc.collect()\n", + " \n", + " # 2. Extract context activations\n", + " internalize_messages = [\n", + " [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"video\", \"path\": example[\"video_path\"]},\n", + " {\"type\": \"text\", \"text\": train_args.internalization_prompt}\n", + " ]\n", + " }\n", + " ]\n", " ]\n", - " ]\n", - " \n", - " vlm_inputs = prepare_smolvlm_inputs(\n", - " processor,\n", - " internalize_messages,\n", - " device,\n", - " video_fps=train_args.video_fps,\n", - " max_frames=train_args.max_frames,\n", - " video_size_longest_edge=train_args.video_size_longest_edge,\n", - " video_load_backend=train_args.video_load_backend\n", - " )\n", + " \n", + " vlm_inputs = prepare_smolvlm_inputs(\n", + " processor,\n", + " internalize_messages,\n", + " device,\n", + " video_fps=train_args.video_fps,\n", + " max_frames=train_args.max_frames,\n", + " video_size_longest_edge=train_args.video_size_longest_edge,\n", + " video_load_backend=train_args.video_load_backend\n", + " )\n", "\n", - " ctx_features, ctx_attn_mask, ctx_position_ids = extract_l2l_fused_text_features(\n", - " raw_model,\n", - " vlm_inputs,\n", - " num_target_layers=model.hypernet.n_layers\n", - " )\n", + " ctx_features, ctx_attn_mask, ctx_position_ids = extract_l2l_fused_text_features(\n", + " raw_model,\n", + " vlm_inputs,\n", + " num_target_layers=model.hypernet.n_layers\n", + " )\n", "\n", - " # 3. Generate dynamic LoRA weights\n", - " generated_loras, _ = model.generate_weights(\n", - " ctx_ids=None,\n", - " ctx_features=ctx_features,\n", - " ctx_attn_mask=ctx_attn_mask,\n", - " ctx_position_ids=ctx_position_ids\n", - " )\n", + " # 3. Generate dynamic LoRA weights\n", + " generated_loras, _ = model.generate_weights(\n", + " ctx_ids=None,\n", + " ctx_features=ctx_features,\n", + " ctx_attn_mask=ctx_attn_mask,\n", + " ctx_position_ids=ctx_position_ids\n", + " )\n", "\n", - " generated_loras = combine_lora(\n", - " generated_loras,\n", - " torch.ones(1, dtype=torch.int32, device=device),\n", - " lora_bias=model.hypernet.get_head_bias() if model.hypernet.config.use_bias else None\n", - " )\n", - " \n", - " # 4. Patch base layers and inject generated weights\n", - " model.patch_lora_forward()\n", - " apply_lora_to_layers(\n", - " model.base_model,\n", - " model.hypernet.layer_indices,\n", - " generated_loras,\n", - " torch.ones(1, dtype=torch.int32, device=device),\n", - " position_ids=None\n", - " )\n", - " \n", - " # 5. Run text-only Q&A inference\n", - " prompt_ids = tokenizer.apply_chat_template(\n", - " [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": [{\"type\": \"text\", \"text\": example[\"prompt\"]}]\n", - " }\n", - " ],\n", - " tokenize=True,\n", - " add_generation_prompt=True,\n", - " return_tensors=\"pt\"\n", - " ).to(device)\n", + " generated_loras = combine_lora(\n", + " generated_loras,\n", + " torch.ones(1, dtype=torch.int32, device=device),\n", + " lora_bias=model.hypernet.get_head_bias() if model.hypernet.config.use_bias else None\n", + " )\n", + " \n", + " # 4. Patch base layers and inject generated weights\n", + " model.patch_lora_forward()\n", + " apply_lora_to_layers(\n", + " model.base_model,\n", + " model.hypernet.layer_indices,\n", + " generated_loras,\n", + " torch.ones(1, dtype=torch.int32, device=device),\n", + " position_ids=None\n", + " )\n", + " \n", + " # 5. Run text-only Q&A inference\n", + " prompt_ids = tokenizer.apply_chat_template(\n", + " [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [{\"type\": \"text\", \"text\": example[\"prompt\"]}]\n", + " }\n", + " ],\n", + " tokenize=True,\n", + " add_generation_prompt=True,\n", + " return_tensors=\"pt\"\n", + " ).to(device)\n", "\n", - " if hasattr(prompt_ids, \"keys\"):\n", - " input_ids = prompt_ids[\"input_ids\"]\n", - " attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n", - " else:\n", - " input_ids = prompt_ids\n", - " attention_mask = torch.ones_like(input_ids)\n", + " if hasattr(prompt_ids, \"keys\"):\n", + " input_ids = prompt_ids[\"input_ids\"]\n", + " attention_mask = prompt_ids.get(\"attention_mask\", torch.ones_like(input_ids))\n", + " else:\n", + " input_ids = prompt_ids\n", + " attention_mask = torch.ones_like(input_ids)\n", "\n", - " generated_ids = model.base_model.generate(\n", - " input_ids=input_ids,\n", - " attention_mask=attention_mask,\n", - " max_new_tokens=train_args.generation_max_new_tokens,\n", - " do_sample=False,\n", - " pad_token_id=tokenizer.pad_token_id,\n", - " eos_token_id=tokenizer.eos_token_id\n", - " )\n", + " generated_ids = model.base_model.generate(\n", + " input_ids=input_ids,\n", + " attention_mask=attention_mask,\n", + " max_new_tokens=train_args.generation_max_new_tokens,\n", + " do_sample=False,\n", + " pad_token_id=tokenizer.pad_token_id,\n", + " eos_token_id=tokenizer.eos_token_id\n", + " )\n", "\n", - " video2lora_pred = tokenizer.decode(\n", - " generated_ids[0][input_ids.shape[1]:],\n", - " skip_special_tokens=True\n", - " ).strip()\n", - " \n", - " video2lora_predictions.append(video2lora_pred)\n", + " video2lora_pred = tokenizer.decode(\n", + " generated_ids[0][input_ids.shape[1]:],\n", + " skip_special_tokens=True\n", + " ).strip()\n", + " \n", + " video2lora_predictions.append(video2lora_pred)\n", + " \n", + " # Clean up temporary GPU tensors for this step\n", + " del vlm_inputs, ctx_features, ctx_attn_mask, ctx_position_ids, generated_loras, prompt_ids, input_ids, attention_mask, generated_ids\n", + " model.reset()\n", + " torch.cuda.empty_cache()\n", + " gc.collect()\n", "\n", "# Reset model hooks back to clean base state\n", "model.reset()\n", + "torch.cuda.empty_cache()\n", + "gc.collect()\n", "\n", "print(\"\\nVideo2LoRA internalization and inferences completed successfully.\")" ]