#!/usr/bin/env python3

import os
import json
import sys

# Add the venv to path so we can import sam2 properly
sys.path.insert(0, '/home/mike/comfyui/venv/lib/python3.13/site-packages')

# This mimics exactly what happens in edit_api.py _load_sam2 function
CONFIG_PATH = "/home/mike/dev/qwen-image-edit-rapid-aio-nsfw-v23/tour-comfy/config.json"

print(f"Testing SAM2 loading with exact code from edit_api.py")
print(f"CONFIG_PATH: {CONFIG_PATH}")

try:
    # This is the exact code from _load_sam2 function
    with open(CONFIG_PATH) as f:
        conf = json.load(f)
    
    print("Config loaded successfully")
    print("sam2_checkpoint:", conf.get("sam2_checkpoint"))
    print("sam2_config:", conf.get("sam2_config"))
    
    # This is the exact code from _load_sam2 function
    ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
    cfg  = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
    
    print(f"Checkpoint path: {ckpt}")
    print(f"Config path: {cfg}")
    print(f"Checkpoint exists: {os.path.exists(ckpt)}")
    
    # Now try to import and load SAM2 exactly as edit_api.py does
    from sam2.build_sam import build_sam2
    from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
    
    if not os.path.exists(ckpt):
        raise FileNotFoundError(f"SAM2 checkpoint not found: {ckpt}")
        
    print("Loading SAM2 model...")
    model = build_sam2(cfg, ckpt, device="cuda")
    print("SUCCESS: SAM2 model loaded successfully!")
    
except Exception as e:
    print(f"ERROR: {e}")
    import traceback
    traceback.print_exc()