#!/usr/bin/env python3
"""
Quick test to verify our custom DFM_Qwen2VL model is loaded correctly
"""
import sys

model_dir = "/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_DFM"

print("=" * 80)
print("Testing DFM_Qwen2VL Model Loading")
print("=" * 80)

print("\n1. Loading config...")
from transformers import AutoConfig
config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
print(f"   ✓ Config loaded: {type(config).__name__}")
print(f"   Model type: {config.model_type}")
print(f"   Base model: {getattr(config, 'base_model_name_or_path', 'N/A')}")

print("\n2. Loading model...")
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    trust_remote_code=True,
    device_map="cpu"
)
print(f"   ✓ Model loaded: {type(model).__name__}")

print("\n3. Checking model components...")
print(f"   Has visual: {hasattr(model, 'visual')}")
print(f"   Has mm_projector: {hasattr(model, 'mm_projector')}")
print(f"   Has model (DFM LLM): {hasattr(model, 'model')}")

if hasattr(model, 'visual'):
    print(f"   Visual parameters: {sum(p.numel() for p in model.visual.parameters()):,}")
    print(f"   Visual frozen: {not any(p.requires_grad for p in model.visual.parameters())}")

if hasattr(model, 'mm_projector') and model.mm_projector is not None:
    print(f"   MM Projector parameters: {sum(p.numel() for p in model.mm_projector.parameters()):,}")
    print(f"   MM Projector frozen: {not any(p.requires_grad for p in model.mm_projector.parameters())}")

if hasattr(model, 'model'):
    print(f"   DFM LLM parameters: {sum(p.numel() for p in model.model.parameters()):,}")

print("\n4. Checking DFM-specific attributes...")
print(f"   Has bd_size: {hasattr(model.config, 'bd_size')}")
if hasattr(model.config, 'bd_size'):
    print(f"   bd_size: {model.config.bd_size}")
print(f"   Has mdm_loss: {hasattr(model.config, 'mdm_loss')}")
if hasattr(model.config, 'mdm_loss'):
    print(f"   mdm_loss: {model.config.mdm_loss}")

print("\n" + "=" * 80)
print("✓ All checks passed! DFM_Qwen2VL model is properly configured.")
print("=" * 80)

