"""Read more of model_management.py — find how --novram affects GPU compute."""
import paramiko

k = paramiko.Ed25519Key.from_private_key_file(r'C:\Users\fabia\.ssh\id_ed25519')
c = paramiko.SSHClient()
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
c.connect('192.168.178.150', username='fabian', pkey=k, timeout=10)
sftp = c.open_sftp()

with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    mgmt = f.read().decode()

lines = mgmt.split('\n')

# Find all functions related to device/offload
print("=== KEY FUNCTIONS ===")
for keyword in ['unet_offload_device', 'unet_device', 'NO_VRAM', 'should_use', 'get_torch_device', 
                'def text_encoder_device', 'def text_encoder_offload', 'VRAMState']:
    for i, line in enumerate(lines):
        if keyword in line and ('def ' in line or 'class ' in line or '=' in line[:50]):
            print(f"  L{i+1}: {line.strip()[:100]}")

# Show VRAMState enum
print("\n=== VRAMState ===")
for i, line in enumerate(lines):
    if 'class VRAMState' in line or (i > 0 and 'VRAMState' in lines[i-1] and 'class' in lines[i-1]):
        for j in range(i, min(i+15, len(lines))):
            print(f"  {j+1}: {lines[j]}")
        break

# Show unet_offload_device
print("\n=== unet_offload_device ===")
for i, line in enumerate(lines):
    if 'def unet_offload_device' in line:
        for j in range(max(0,i-2), min(i+15, len(lines))):
            print(f"  {j+1}: {lines[j]}")

# Show text_encoder functions
print("\n=== text_encoder_device ===")
for i, line in enumerate(lines):
    if 'def text_encoder_device' in line:
        for j in range(max(0,i-2), min(i+12, len(lines))):
            print(f"  {j+1}: {lines[j]}")

print("\n=== text_encoder_offload_device ===")
for i, line in enumerate(lines):
    if 'def text_encoder_offload_device' in line:
        for j in range(max(0,i-2), min(i+12, len(lines))):
            print(f"  {j+1}: {lines[j]}")

# Show how NO_VRAM is used in loading logic
print("\n=== NO_VRAM usage in model loading ===")
for i, line in enumerate(lines):
    if 'NO_VRAM' in line:
        print(f"  L{i+1}: {line.strip()[:120]}")

# Show the VRAM state setting logic
print("\n=== vram_state assignment ===")
for i, line in enumerate(lines):
    if 'vram_state' in line and ('=' in line) and 'VRAMState' in line:
        print(f"  L{i+1}: {line.strip()[:120]}")

sftp.close()
c.close()