77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
"""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.")
|