"""Reconnect, check patches, fix missing ones, restart."""
import paramiko, time, json

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()

def sh(cmd, t=30):
    ch = c.get_transport().open_session()
    ch.settimeout(t)
    ch.exec_command('/bin/bash -l -c ' + "'" + cmd.replace("'", "'\\''") + "'")
    o = b""
    while True:
        try:
            d = ch.recv(65536)
            if not d: break
            o += d
        except: break
    ch.close()
    return o.decode(errors='replace').strip()

# Read current state
with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    code = f.read().decode()
lines = code.split('\n')

# Show ALL offload functions
for fname in ['unet_offload_device', 'vae_offload_device', 'text_encoder_offload_device']:
    for i, line in enumerate(lines):
        if f'def {fname}' in line:
            print(f"\n=== {fname} (L{i+1}) ===")
            for j in range(i, min(i+10, len(lines))):
                print(f"  {j+1}: {lines[j]}")
            break

# Also show unet_inital_load_device  
for i, line in enumerate(lines):
    if 'def unet_inital_load_device' in line:
        print(f"\n=== unet_inital_load_device (L{i+1}) ===")
        for j in range(i, min(i+10, len(lines))):
            print(f"  {j+1}: {lines[j]}")
        break

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

sftp.close()
c.close()
print("\nDone reading.")