"""Patch ComfyUI: Force VAE to GPU even in --novram mode. Restart and test."""
import paramiko, time, json, textwrap

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=15)
sftp = c.open_sftp()

def sh(cmd, timeout=60):
    chan = c.get_transport().open_session()
    chan.settimeout(timeout)
    chan.exec_command('/bin/bash -l -c ' + "'" + cmd.replace("'", "'\\''") + "'")
    out = b""
    while True:
        try:
            chunk = chan.recv(65536)
            if not chunk: break
            out += chunk
        except: break
    chan.close()
    return out.decode(errors='replace').strip()

# =============================================
# STEP 1: Kill ComfyUI
# =============================================
print("1) Kill ComfyUI")
sh('pkill -9 -f "python3.*main.py" 2>/dev/null; sleep 2')

# =============================================
# STEP 2: Read and patch model_management.py
# =============================================
print("2) Patch model_management.py — force VAE to GPU")

# First, read the file to understand the structure
with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    mgmt = f.read().decode()

print(f"   File size: {len(mgmt)} bytes")

# Find the vae_offload_device function
# In ComfyUI, with NO_VRAM, vae_offload_device() returns CPU
# We need to make it return GPU instead

# Also find vae_dtype — it's set to float32 by default, we want float16

# Let's search for relevant functions
for i, line in enumerate(mgmt.split('\n')):
    if 'def vae_offload_device' in line or 'def vae_dtype' in line or 'def vae_device' in line:
        print(f"   Line {i+1}: {line.strip()}")

# Also check what functions exist
found = []
for i, line in enumerate(mgmt.split('\n')):
    if line.startswith('def ') or (line.startswith('    ') and 'def ' in line[:12]):
        if 'vae' in line.lower():
            found.append((i+1, line.strip()))
for ln, l in found:
    print(f"   L{ln}: {l}")

# Let's read the specific area around these functions
lines = mgmt.split('\n')

# Find and show context around vae functions
for keyword in ['vae_offload_device', 'vae_dtype', 'vae_device']:
    for i, line in enumerate(lines):
        if f'def {keyword}' in line:
            start = max(0, i-2)
            end = min(len(lines), i+15)
            print(f"\n   --- {keyword} (L{i+1}) ---")
            for j in range(start, end):
                print(f"   {j+1:>5}: {lines[j]}")

sftp.close()
c.close()
print("\n   Reading complete. Will patch next.")