317 lines
12 KiB
Python
317 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Fix: Add MIOpen fast-find, pre-warm GPU, restart ComfyUI, generate.
|
|
All GPU ops confirmed working. The hang is likely cold MIOpen kernel cache.
|
|
Single SSH connection.
|
|
"""
|
|
import paramiko, json, time, sys
|
|
|
|
ssh = paramiko.SSHClient()
|
|
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
for attempt in range(5):
|
|
try:
|
|
ssh.connect('192.168.178.150', username='fabian',
|
|
key_filename=r'C:\Users\fabia\.ssh\id_ed25519', timeout=10)
|
|
break
|
|
except Exception as e:
|
|
print(f" SSH attempt {attempt+1}/5: {e}")
|
|
time.sleep(10)
|
|
else:
|
|
print("FATAL: Cannot connect"); sys.exit(1)
|
|
|
|
def run(cmd, timeout=300):
|
|
_, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
|
|
out = stdout.read().decode()
|
|
err = stderr.read().decode()
|
|
return out, err
|
|
|
|
try:
|
|
# 1. Kill any lingering ComfyUI
|
|
print("=== Kill any ComfyUI ===")
|
|
run("pkill -9 -f 'python3 main.py' 2>/dev/null; sleep 1")
|
|
print(" Done")
|
|
|
|
# 2. Pre-warm MIOpen kernel cache with DiT-like operations
|
|
print("\n=== Pre-warming MIOpen kernel cache ===")
|
|
print(" This compiles HIP kernels that Z-Image-Turbo will need.")
|
|
print(" First run after reboot is slow (kernel compilation)...")
|
|
|
|
warmup_code = r"""
|
|
import torch, torch.nn as nn, time
|
|
|
|
# Simulate Z-Image-Turbo DiT operations
|
|
device = 'cuda'
|
|
dtype = torch.float16
|
|
|
|
print("Warming up HIP kernels for DiT inference...")
|
|
t0 = time.time()
|
|
|
|
# 1. Linear layers (DiT blocks)
|
|
print(" Linear layers...", end=" ", flush=True)
|
|
for size in [(1024,1024), (4096,1024), (1024,4096)]:
|
|
l = nn.Linear(*size).to(device, dtype)
|
|
x = torch.randn(1, 64, size[0], device=device, dtype=dtype)
|
|
y = l(x)
|
|
del l, x, y
|
|
torch.cuda.synchronize()
|
|
print(f"{time.time()-t0:.1f}s")
|
|
|
|
# 2. Attention (SDPA - the core of DiT)
|
|
print(" Scaled dot-product attention...", end=" ", flush=True)
|
|
t1 = time.time()
|
|
for heads in [8, 16, 24]:
|
|
q = torch.randn(1, heads, 64, 64, device=device, dtype=dtype)
|
|
k = torch.randn(1, heads, 64, 64, device=device, dtype=dtype)
|
|
v = torch.randn(1, heads, 64, 64, device=device, dtype=dtype)
|
|
y = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
del q, k, v, y
|
|
torch.cuda.synchronize()
|
|
print(f"{time.time()-t1:.1f}s")
|
|
|
|
# 3. LayerNorm / RMSNorm
|
|
print(" Normalization layers...", end=" ", flush=True)
|
|
t1 = time.time()
|
|
for dim in [1024, 2048, 4096]:
|
|
ln = nn.LayerNorm(dim).to(device, dtype)
|
|
x = torch.randn(1, 64, dim, device=device, dtype=dtype)
|
|
y = ln(x)
|
|
del ln, x, y
|
|
torch.cuda.synchronize()
|
|
print(f"{time.time()-t1:.1f}s")
|
|
|
|
# 4. Conv2d (VAE-like, but we'll run VAE on CPU)
|
|
print(" Conv2d layers...", end=" ", flush=True)
|
|
t1 = time.time()
|
|
for ch in [64, 128, 256]:
|
|
c = nn.Conv2d(ch, ch, 3, padding=1).to(device, dtype)
|
|
x = torch.randn(1, ch, 32, 32, device=device, dtype=dtype)
|
|
y = c(x)
|
|
del c, x, y
|
|
torch.cuda.synchronize()
|
|
print(f"{time.time()-t1:.1f}s")
|
|
|
|
# 5. Full mini-DiT forward pass simulation
|
|
print(" Mini-DiT forward pass simulation...", end=" ", flush=True)
|
|
t1 = time.time()
|
|
hidden = 1024
|
|
seq_len = 256
|
|
heads = 16
|
|
head_dim = hidden // heads
|
|
# Simulate a DiT block
|
|
x = torch.randn(1, seq_len, hidden, device=device, dtype=dtype)
|
|
norm = nn.LayerNorm(hidden).to(device, dtype)
|
|
qkv = nn.Linear(hidden, hidden*3).to(device, dtype)
|
|
proj = nn.Linear(hidden, hidden).to(device, dtype)
|
|
ff1 = nn.Linear(hidden, hidden*4).to(device, dtype)
|
|
ff2 = nn.Linear(hidden*4, hidden).to(device, dtype)
|
|
for step in range(3):
|
|
h = norm(x)
|
|
q, k, v = qkv(h).chunk(3, dim=-1)
|
|
q = q.view(1, seq_len, heads, head_dim).transpose(1,2)
|
|
k = k.view(1, seq_len, heads, head_dim).transpose(1,2)
|
|
v = v.view(1, seq_len, heads, head_dim).transpose(1,2)
|
|
attn = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
attn = attn.transpose(1,2).contiguous().view(1, seq_len, hidden)
|
|
x = x + proj(attn)
|
|
x = x + ff2(torch.nn.functional.gelu(ff1(norm(x))))
|
|
torch.cuda.synchronize()
|
|
print(f"{time.time()-t1:.1f}s")
|
|
|
|
total = time.time() - t0
|
|
print(f"\nGPU kernel warmup complete in {total:.1f}s")
|
|
print(f"VRAM used: {torch.cuda.memory_allocated()//1048576} MB")
|
|
torch.cuda.empty_cache()
|
|
print(f"VRAM after cleanup: {torch.cuda.memory_allocated()//1048576} MB")
|
|
print("WARMUP_DONE")
|
|
"""
|
|
# Write warmup script
|
|
run(f"cat > /tmp/gpu_warmup.py << 'PYEOF'\n{warmup_code}\nPYEOF")
|
|
|
|
out, err = run("bash -c 'source ~/comfyui-env/bin/activate && "
|
|
"HSA_OVERRIDE_GFX_VERSION=10.1.0 HSA_ENABLE_SDMA=0 "
|
|
"MIOPEN_FIND_MODE=3 MIOPEN_FIND_ENFORCE=3 "
|
|
"python3 /tmp/gpu_warmup.py' 2>&1", timeout=300)
|
|
print(out.strip())
|
|
if 'WARMUP_DONE' not in out:
|
|
print(f" WARNING: Warmup may have failed")
|
|
print(f" STDERR: {err.strip()[:500]}")
|
|
|
|
# 3. Update startup script with MIOpen settings
|
|
print("\n=== Update startup script ===")
|
|
script = """#!/bin/bash
|
|
# BC-250 ComfyUI Launcher — GPU (ROCm) + CPU VAE
|
|
|
|
# GPU identity (Cyan Skillfish gfx1013 -> gfx1010)
|
|
export HSA_OVERRIDE_GFX_VERSION=10.1.0
|
|
export HIP_VISIBLE_DEVICES=0
|
|
export HSA_ENABLE_SDMA=0
|
|
export HSA_TOOLS_LIB=""
|
|
export HSA_TOOLS_REPORT_LOAD_FAILURE=0
|
|
|
|
# MIOpen: fast kernel selection (avoid long auto-tune on first run)
|
|
export MIOPEN_FIND_MODE=3
|
|
export MIOPEN_FIND_ENFORCE=3
|
|
|
|
# Use all 12 CPU cores for CPU-side work (dequant, text encoding)
|
|
export OMP_NUM_THREADS=12
|
|
export MKL_NUM_THREADS=12
|
|
export OPENBLAS_NUM_THREADS=12
|
|
|
|
# Activate venv
|
|
source /home/fabian/comfyui-env/bin/activate
|
|
cd /home/fabian/ComfyUI
|
|
|
|
# --novram: offload models to RAM, send layers to GPU one at a time
|
|
# --force-fp16: fp16 diffusion to halve VRAM usage
|
|
# --cpu-vae: VAE decode on CPU (GPU hangs on full VAE forward pass)
|
|
exec python3 main.py \\
|
|
--listen 0.0.0.0 --port 8188 \\
|
|
--novram \\
|
|
--force-fp16 \\
|
|
--cpu-vae
|
|
"""
|
|
run(f"cat > /home/fabian/start_comfyui.sh << 'HEREDOC_END'\n{script}HEREDOC_END\n"
|
|
f"chmod +x /home/fabian/start_comfyui.sh")
|
|
print(" Written with MIOpen fast-find + --novram --force-fp16 --cpu-vae")
|
|
|
|
# 4. Launch ComfyUI
|
|
print("\n=== Launch ComfyUI ===")
|
|
run("nohup /home/fabian/start_comfyui.sh > /home/fabian/comfyui.log 2>&1 &")
|
|
time.sleep(2)
|
|
|
|
print(" Waiting for server...")
|
|
for i in range(50):
|
|
time.sleep(3)
|
|
out, _ = run("curl -s -o /dev/null -w '%{http_code}' http://localhost:8188/ 2>/dev/null || echo 0")
|
|
if out.strip() == '200':
|
|
print(f" Server ready ({(i+1)*3}s)")
|
|
break
|
|
if i % 5 == 4:
|
|
log, _ = run("tail -2 /home/fabian/comfyui.log 2>/dev/null")
|
|
last = [l.strip() for l in log.strip().split('\n') if l.strip()]
|
|
print(f" [{(i+1)*3}s] ... {last[-1][:80] if last else ''}")
|
|
else:
|
|
print(" Timeout!")
|
|
out, _ = run("tail -40 /home/fabian/comfyui.log")
|
|
print(out)
|
|
sys.exit(1)
|
|
|
|
# 5. Submit workflow
|
|
print("\n=== Submit workflow (512x512, 8 steps) ===")
|
|
workflow = {
|
|
"prompt": {
|
|
"1": {"class_type": "UnetLoaderGGUF",
|
|
"inputs": {"unet_name": "z_image_turbo-Q5_K_S.gguf"}},
|
|
"2": {"class_type": "CLIPLoaderGGUF",
|
|
"inputs": {"clip_name": "Qwen3-4B.i1-Q5_K_S.gguf", "type": "qwen_image"}},
|
|
"3": {"class_type": "VAELoader",
|
|
"inputs": {"vae_name": "ae.safetensors"}},
|
|
"4": {"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "A red fox in a snowy forest, photorealistic", "clip": ["2", 0]}},
|
|
"5": {"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "", "clip": ["2", 0]}},
|
|
"6": {"class_type": "EmptyLatentImage",
|
|
"inputs": {"width": 512, "height": 512, "batch_size": 1}},
|
|
"7": {"class_type": "KSampler",
|
|
"inputs": {"model": ["1", 0], "seed": 42, "steps": 8, "cfg": 1.0,
|
|
"sampler_name": "euler", "scheduler": "simple",
|
|
"positive": ["4", 0], "negative": ["5", 0],
|
|
"latent_image": ["6", 0], "denoise": 1.0}},
|
|
"8": {"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["7", 0], "vae": ["3", 0]}},
|
|
"9": {"class_type": "SaveImage",
|
|
"inputs": {"images": ["8", 0], "filename_prefix": "ZImageTurbo_test"}}
|
|
}
|
|
}
|
|
|
|
wf_json = json.dumps(workflow)
|
|
# Write workflow to file to avoid shell escaping issues
|
|
run(f"cat > /tmp/zimage_wf.json << 'JSONEOF'\n{wf_json}\nJSONEOF")
|
|
out, _ = run("curl -s -X POST http://localhost:8188/prompt "
|
|
"-H 'Content-Type: application/json' "
|
|
"-d @/tmp/zimage_wf.json")
|
|
try:
|
|
resp = json.loads(out.strip())
|
|
if 'error' in resp:
|
|
print(f" API ERROR: {resp['error']}")
|
|
sys.exit(1)
|
|
print(f" Prompt ID: {resp.get('prompt_id')}")
|
|
except:
|
|
print(f" Response: {out.strip()[:500]}")
|
|
|
|
# 6. Monitor — wait up to 15 minutes (first run can be slow due to kernel cache)
|
|
print("\n=== Monitoring generation (GPU kernels may compile on first step) ===")
|
|
last_log = ""
|
|
for i in range(60): # up to 15 minutes
|
|
time.sleep(15)
|
|
|
|
stats, _ = run("bash -c '"
|
|
"PID=$(pgrep -f \"python3 main.py\" | head -1); "
|
|
"if [ -n \"$PID\" ]; then "
|
|
" CPU=$(ps -p $PID -o %cpu --no-headers); "
|
|
" RSS=$(ps -p $PID -o rss --no-headers); "
|
|
" LOAD=$(cat /proc/loadavg | cut -d\" \" -f1); "
|
|
" GPU_T=$(cat /sys/class/drm/card1/device/hwmon/hwmon*/temp1_input 2>/dev/null || echo 0); "
|
|
" echo \"CPU:${CPU}% RSS:$((RSS/1024))M LOAD:${LOAD} GPU:$((GPU_T/1000))C\"; "
|
|
"else echo DEAD; fi'")
|
|
|
|
log, _ = run("tail -12 /home/fabian/comfyui.log 2>/dev/null")
|
|
log_s = log.strip()
|
|
|
|
elapsed = (i+1) * 15
|
|
m, s = divmod(elapsed, 60)
|
|
stats_s = stats.strip()
|
|
print(f" [{m}m{s:02d}s] {stats_s}")
|
|
|
|
# Show new log content
|
|
if log_s != last_log:
|
|
for line in reversed(log_s.split('\n')):
|
|
l = line.strip()
|
|
if l and not l.startswith('FETCH') and not l.startswith('[DEPRECATION') and not l.startswith('[ComfyUI-Manager]'):
|
|
print(f" LOG: {l[:120]}")
|
|
break
|
|
last_log = log_s
|
|
|
|
if 'DEAD' in stats_s:
|
|
print("\n PROCESS DIED!")
|
|
out, _ = run("tail -60 /home/fabian/comfyui.log")
|
|
print(out)
|
|
break
|
|
|
|
if 'Prompt executed in' in log_s:
|
|
print(f"\n SUCCESS! Image generated!")
|
|
out, _ = run("tail -25 /home/fabian/comfyui.log")
|
|
print(out)
|
|
break
|
|
|
|
if 'Traceback' in log_s or 'RuntimeError' in log_s:
|
|
print("\n ERROR detected!")
|
|
out, _ = run("tail -60 /home/fabian/comfyui.log")
|
|
print(out)
|
|
break
|
|
|
|
# 7. Check output
|
|
print("\n=== Output files ===")
|
|
out, _ = run("ls -lah ~/ComfyUI/output/ 2>/dev/null")
|
|
print(out.strip())
|
|
|
|
# 8. Check history
|
|
out, _ = run("curl -s http://localhost:8188/history 2>/dev/null")
|
|
try:
|
|
h = json.loads(out)
|
|
for pid, info in h.items():
|
|
status = info.get('status', {})
|
|
outputs = info.get('outputs', {})
|
|
print(f"\n Prompt {pid[:12]}...: status={status}")
|
|
if outputs:
|
|
for nid, nout in outputs.items():
|
|
if isinstance(nout, dict) and 'images' in nout:
|
|
for img in nout['images']:
|
|
print(f" Image: {img.get('filename', 'unknown')}")
|
|
except:
|
|
pass
|
|
|
|
finally:
|
|
ssh.close()
|
|
print("\nSSH connection closed.")
|