#!/usr/bin/env python3 # ================================================================================ # LUMEN LIGHTNING WORKER -- Wan 2.1 14B (GGUF) + CausVid 4-step + TeaCache # ================================================================================ # Runs inside a Lightning AI Studio (GPU: T4 or L4). Serves a plain REST API on # port 8000 -- NO gradio, NO tunnel. Then expose port 8000 with the Studio's # "Ports" plugin to get a stable public URL, and paste that into Lumen. # # WHY this shape: Lightning Studios have PERSISTENT DISK, so ComfyUI + the 10GB # model download ONCE and stay. The clip is returned as base64 over the stable # Lightning URL (no trycloudflare truncation), and /result reports the exact byte # count so the client can verify the file is complete (moov atom) before accepting. # # RUN: python lightning_worker.py # (first run ~10-15 min: installs ComfyUI + pulls ~16GB of weights, cached # after. Then it prints "SERVING on :8000" -> expose the port -> paste URL.) # # Endpoints (same contract Modal used): # GET / -> {"ok": true, ...} health # POST /submit {prompt, steps} -> {"call_id": "..."} starts a render # GET /result?call_id=... -> {"status": "running"|"done"|"error", # "bytes": N, "video_b64": "..."} (done -> the mp4) # ================================================================================ import os, sys, subprocess, time, threading, json, base64, urllib.request, urllib.error, shutil HOME = os.path.expanduser("~") ROOT = os.path.join(HOME, "lumen_worker") COMFY_DIR = os.path.join(ROOT, "ComfyUI") QUANT = os.environ.get("LUMEN_QUANT", "Q5_K_S") # fits T4 16GB; Q4_K_M if tight PORT = int(os.environ.get("LUMEN_PORT", "8000")) COMFY_PORT = 8188 WIDTH, HEIGHT = 832, 480 FRAMES, FPS = 81, 16 STEPS = int(os.environ.get("LUMEN_STEPS", "4")) CFG = 1.0 SHIFT = 8.0 TEACACHE_THRESH = 0.20 CAUSVID_STR = 1.0 NEG = "worst quality, low quality, blurry, jittery, distorted, watermark, text, deformed" _PHYS = [ (("pendulum","swing","swinging"),"pendulum motion accelerating through the lowest point and slowing at each extreme"), (("orbit","orbiting","planet","moon","celestial"),"smooth orbital motion, faster when nearer, conserving angular momentum"), (("throw","thrown","launch","arc","parabola","ballistic"),"a smooth parabolic arc, rising then accelerating downward under gravity"), (("bounce","bouncing","bounces"),"realistic bouncing with energy loss, each rebound lower, gravity-accelerated between bounces"), (("collision","collide","impact","crash","momentum"),"a momentum-conserving collision, motion transferring realistically on impact"), (("spring","oscillat","vibrat","harmonic"),"damped harmonic oscillation, smooth decaying back-and-forth motion"), (("fall","falling","drop","dropping","plummet"),"a gravity-accelerated fall, speed increasing as it descends"), (("splash","waterfall","pour","fluid","flowing water"),"gravity-driven fluid flow with realistic acceleration and splashing"), (("explosion","explode","blast","debris"),"explosive expansion followed by gravity-driven debris arcs"), (("smoke","fire","flame","steam"),"buoyant rising motion with turbulent, natural flow"), (("run","running","walk","walking","jump","sprint"),"natural human biomechanics with believable weight, momentum and balance"), (("car","driving","wheel","motorcycle","truck","vehicle"),"realistic vehicle dynamics, wheels rotating in sync with speed"), (("flag","cloth","fabric","hair","cape","curtain"),"cloth physics, natural billowing and settling"), ] def apply_physics(prompt): p=(prompt or "").lower(); hits=[] for keys,cue in _PHYS: if any(k in p for k in keys) and cue not in hits: hits.append(cue) if len(hits)>=2: break base=(prompt or "a cinematic scene").strip() return ("%s, %s"%(base, ", ".join(hits))) if hits else base def sh(cmd): print("+ "+cmd, flush=True); subprocess.run(cmd, shell=True, check=False) # ---- 1. install ComfyUI + custom nodes (idempotent; persists on Studio disk) --- def setup(): os.makedirs(ROOT, exist_ok=True) if not os.path.isdir(COMFY_DIR): sh("git clone --depth=1 https://github.com/comfyanonymous/ComfyUI '%s'"%COMFY_DIR) sh("%s -m pip install -q -r '%s/requirements.txt'"%(sys.executable, COMFY_DIR)) nodes = { "ComfyUI-GGUF":"https://github.com/city96/ComfyUI-GGUF", "ComfyUI-VideoHelperSuite":"https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite", "ComfyUI-TeaCache":"https://github.com/welltop-cn/ComfyUI-TeaCache", "ComfyUI-MagCache":"https://github.com/Zehong-Ma/ComfyUI-MagCache", } cn = os.path.join(COMFY_DIR,"custom_nodes") for name,url in nodes.items(): d=os.path.join(cn,name) if not os.path.isdir(d): sh("git clone --depth=1 %s '%s'"%(url,d)) req=os.path.join(d,"requirements.txt") if os.path.exists(req): sh("%s -m pip install -q -r '%s'"%(sys.executable,req)) # FIX: TeaCache's line-6 `from comfy.ldm.lightricks.model import precompute_freqs_cis` # is LTX-Video-ONLY (Wan never uses it), but a newer ComfyUI dropped that symbol, so the # top-level import kills the WHOLE node -> no TeaCache. Wrap it in try/except. Idempotent; # runs every launch, so it also repairs an already-cloned node on the next restart. _tc=os.path.join(cn,"ComfyUI-TeaCache","nodes.py") try: _s=open(_tc,encoding="utf-8").read() _bad="from comfy.ldm.lightricks.model import precompute_freqs_cis" if _bad in _s and "precompute_freqs_cis = None" not in _s: open(_tc,"w",encoding="utf-8").write(_s.replace(_bad, "try:\n "+_bad+"\nexcept Exception:\n precompute_freqs_cis = None # LTXV-only; Wan unaffected")) print(" [patch] TeaCache import made optional -> node will load (Wan-safe)",flush=True) except Exception: pass # huggingface_hub: NOT pinned to 0.25.2 anymore (that was for the old gradio worker; this # REST worker doesn't use gradio). The Studio's transformers 5.14 needs hf_hub>=1.5 or it # fails with "cannot import name 'is_offline_mode'". opencv<5: opencv 5.x demands numpy>=2, # which fights scipy's numpy<2 need; 4.x works with numpy<2. hf_hub_download works on 1.x. sh("%s -m pip install -q gguf 'fastapi[standard]' uvicorn 'huggingface_hub>=1.5,<2' imageio-ffmpeg 'opencv-python-headless<5'"%sys.executable) # FIX (numpy ABI): the installs above pull NumPy 2.x, whose C ABI breaks the Studio's # prebuilt scipy that ComfyUI imports -> "numpy.core.multiarray failed to import" -> # ComfyUI can't even start. Pin NumPy back to <2. LAST, so it wins over everything else. # ComfyUI launches as a fresh subprocess so it picks this up with no kernel restart. sh("%s -m pip install -q 'numpy<2'"%sys.executable) # FIX (torchaudio ABI): the env's torchaudio was built for an OLDER torch than the # installed 2.8.0+cu128 -> "undefined symbol: torch_library_impl" -> ComfyUI (which # imports torchaudio via the LTX audio VAE we don't use) crashes at startup. Reinstall # torchaudio matching torch 2.8.0; --no-deps so it never disturbs torch itself. sh("%s -m pip install -q --force-reinstall --no-deps 'torchaudio==2.8.0' --index-url https://download.pytorch.org/whl/cu128"%sys.executable) # ---- 2. download weights (idempotent) ---------------------------------------- MODELS = {} def download(): from huggingface_hub import hf_hub_download M=os.path.join(COMFY_DIR,"models") for sub in ("unet","text_encoders","vae","loras"): os.makedirs(os.path.join(M,sub),exist_ok=True) def dl(repo,cands,sub,label,required=True): for fn in cands: dst=os.path.join(M,sub,os.path.basename(fn)) if os.path.exists(dst): print(" %s: %s (cached)"%(label,os.path.basename(fn)),flush=True); return os.path.basename(fn) last=None for fn in cands: try: p=hf_hub_download(repo,fn,local_dir=os.path.join(M,sub)) nm=os.path.basename(p) dstf=os.path.join(M,sub,nm) if p!=dstf and not os.path.exists(dstf): try: shutil.move(p,dstf) except Exception: pass print(" %s: %s (downloaded)"%(label,nm),flush=True); return nm except Exception as e: last=e print(" !! %s missing: %s"%(label,repr(last)[:80]),flush=True) if required: raise RuntimeError("required weight missing: "+label) return None print("[1/3] weights (first run ~16GB, cached after)...",flush=True) MODELS["t2v"]=dl("city96/Wan2.1-T2V-14B-gguf",["wan2.1-t2v-14b-%s.gguf"%QUANT],"unet","t2v") MODELS["txt"]=dl("Comfy-Org/Wan_2.1_ComfyUI_repackaged", ["split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors"],"text_encoders","text") MODELS["vae"]=dl("Comfy-Org/Wan_2.1_ComfyUI_repackaged", ["split_files/vae/wan_2.1_vae.safetensors"],"vae","vae") MODELS["causvid"]=dl("Kijai/WanVideo_comfy", ["Wan21_CausVid_14B_T2V_lora_rank32_v2.safetensors", "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"],"loras","causvid",required=False) # ---- 3. start ComfyUI -------------------------------------------------------- TEACACHE_NODE=[None]; MAGCACHE_NODE=[None] def start_comfy(): print("[2/3] starting ComfyUI...",flush=True) # GPU/torch sanity -- prints immediately whether the T4 is actually visible to torch try: import torch print(" torch %s | CUDA=%s | %s" % (torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU-ONLY (!)"), flush=True) except Exception as e: print(" torch check failed: %s" % repr(e)[:120], flush=True) # kill any ComfyUI left over from earlier attempts (frees port 8188 + GPU memory) subprocess.run("pkill -9 -f 'ComfyUI/main.py' 2>/dev/null; sleep 1; true", shell=True) logpath = os.path.join(ROOT, "comfy.log") log = open(logpath, "w") proc = subprocess.Popen([sys.executable, os.path.join(COMFY_DIR, "main.py"), "--listen", "127.0.0.1", "--port", str(COMFY_PORT), "--disable-metadata", "--disable-auto-launch"], cwd=COMFY_DIR, stdout=log, stderr=subprocess.STDOUT) up = False for i in range(150): time.sleep(2) if proc.poll() is not None: # ComfyUI process EXITED -> it crashed; stop waiting print(" ComfyUI exited early (code %s)" % proc.returncode, flush=True) break try: if urllib.request.urlopen("http://127.0.0.1:%d/system_stats" % COMFY_PORT, timeout=3).status == 200: up = True; print(" ComfyUI up in %ds" % ((i + 1) * 2), flush=True); break except Exception: pass if not up: # print the ACTUAL error straight here so we don't have to open comfy.log separately try: log.flush() except Exception: pass try: tail = open(logpath, encoding="utf-8", errors="replace").read()[-3500:] except Exception: tail = "(could not read comfy.log)" print("\n" + "="*26 + " ComfyUI FAILED -- tail of comfy.log " + "="*26, flush=True) print(tail, flush=True) print("="*88, flush=True) raise RuntimeError("ComfyUI did not start -- read the comfy.log dump printed just above") try: info=json.loads(urllib.request.urlopen("http://127.0.0.1:%d/object_info"%COMFY_PORT,timeout=15).read()) for n in ("TeaCache","TeaCacheForVidGen","WanVideoTeaCache","ApplyTeaCache"): if n in info: TEACACHE_NODE[0]=n; break if not TEACACHE_NODE[0]: for n in ("MagCache","MagCacheForVideo","ApplyMagCache","WanMagCache"): if n in info: MAGCACHE_NODE[0]=n; break print(" cache: TeaCache=%s MagCache=%s"%(TEACACHE_NODE[0],MAGCACHE_NODE[0]),flush=True) except Exception as e: print(" cache probe failed (%s)"%repr(e)[:50],flush=True) print("[3/3] ComfyUI ready.",flush=True) # ---- 4. graph + render ------------------------------------------------------- def build(prompt,seed,prefix,cache=True): G={} G["1"]={"class_type":"UnetLoaderGGUF","inputs":{"unet_name":MODELS["t2v"]}}; model=["1",0] if MODELS.get("causvid"): G["1c"]={"class_type":"LoraLoaderModelOnly","inputs":{"model":model,"lora_name":MODELS["causvid"],"strength_model":CAUSVID_STR}}; model=["1c",0] G["6"]={"class_type":"ModelSamplingSD3","inputs":{"model":model,"shift":SHIFT}}; model=["6",0] if cache and TEACACHE_NODE[0]: # TeaCache REQUIRES model_type (enum) -- for our model it's wan2.1_t2v_14B. G["6t"]={"class_type":TEACACHE_NODE[0],"inputs":{"model":model,"model_type":"wan2.1_t2v_14B", "rel_l1_thresh":TEACACHE_THRESH,"start_percent":0.0,"end_percent":1.0,"cache_device":"cuda"}}; model=["6t",0] elif cache and MAGCACHE_NODE[0]: G["6c"]={"class_type":MAGCACHE_NODE[0],"inputs":{"model":model,"threshold":0.12}}; model=["6c",0] G["2"]={"class_type":"CLIPLoader","inputs":{"clip_name":MODELS["txt"],"type":"wan"}} G["4"]={"class_type":"CLIPTextEncode","inputs":{"clip":["2",0],"text":prompt}} G["5"]={"class_type":"CLIPTextEncode","inputs":{"clip":["2",0],"text":NEG}} G["3"]={"class_type":"VAELoader","inputs":{"vae_name":MODELS["vae"]}} G["7"]={"class_type":"EmptyHunyuanLatentVideo","inputs":{"width":WIDTH,"height":HEIGHT,"length":FRAMES,"batch_size":1}} G["8"]={"class_type":"KSampler","inputs":{"model":model,"positive":["4",0],"negative":["5",0],"latent_image":["7",0],"seed":seed,"steps":STEPS,"cfg":CFG,"sampler_name":"euler","scheduler":"simple","denoise":1.0}} G["9"]={"class_type":"VAEDecode","inputs":{"samples":["8",0],"vae":["3",0]}} G["10"]={"class_type":"VHS_VideoCombine","inputs":{"images":["9",0],"frame_rate":FPS,"loop_count":0,"filename_prefix":prefix,"format":"video/h264-mp4","pix_fmt":"yuv420p","crf":18,"pingpong":False,"save_output":True}} return G def run_graph(wf,timeout=1500): body=json.dumps({"prompt":wf}).encode() try: r=urllib.request.urlopen(urllib.request.Request("http://127.0.0.1:%d/prompt"%COMFY_PORT,data=body,headers={"Content-Type":"application/json"},method="POST"),timeout=15) pid=json.loads(r.read()).get("prompt_id") except urllib.error.HTTPError as he: raise RuntimeError("ComfyUI rejected: "+he.read().decode("utf-8","replace")[:400]) t0=time.time() while time.time()-t0=2 and m[0]=="execution_error": d=m[1] or {} detail=" node %s (%s): %s"%(d.get("node_id"),d.get("node_type"),str(d.get("exception_message"))[:220]) break raise RuntimeError("ComfyUI error:"+(detail or " see comfy.log")) except urllib.error.URLError: continue raise RuntimeError("render timed out") # ---- 5. in-process job queue + FastAPI --------------------------------------- _JOBS={} def _spawn(prompt,steps): jid="g%d"%int(time.time()*1000) _JOBS[jid]={"status":"running","data":None,"detail":""} def _run(): try: pr=apply_physics(prompt); stamp=int(time.time()*1000) try: path=run_graph(build(pr[:500],stamp%2**31,"wan14b_%d"%stamp,cache=True)) except RuntimeError as e: m=str(e).lower() if "rejected" in m or "comfyui error" in m: # cache node failed (validation OR runtime) -> render without it print(">>> render failed WITH cache: %s -> retrying WITHOUT cache" % repr(e)[:160], flush=True) path=run_graph(build(pr[:500],stamp%2**31,"wan14b_%d"%stamp,cache=False)) else: raise with open(path,"rb") as f: data=f.read() _JOBS[jid]={"status":"done","data":data,"detail":os.path.basename(path)} print(">>> rendered %s (%d bytes)"%(os.path.basename(path),len(data)),flush=True) except Exception as e: _JOBS[jid]={"status":"error","data":None,"detail":repr(e)[:200]} print(">>> render FAILED: %s"%repr(e)[:200],flush=True) threading.Thread(target=_run,daemon=True).start() return jid # ---- 5b. VIDEO-TO-VIDEO (2026-09-02, FIRST DRAFT — needs a live test) --------- # Same model, but the KSampler starts from the ENCODED input video latent instead of # empty noise, with denoise < 1.0 = "keep the motion/structure, restyle by `strength`". # >>> UNCERTAIN: (a) VHS_LoadVideoPath input keys across VideoHelperSuite versions; # (b) whether ComfyUI's VAEEncode does Wan's 3D (temporal) encode so the latent # length matches the sampler. If it errors, the worker self-reports the node. def build_v2v(prompt,video_path,seed,prefix,strength,cache=True): G={} G["1"]={"class_type":"UnetLoaderGGUF","inputs":{"unet_name":MODELS["t2v"]}}; model=["1",0] if MODELS.get("causvid"): G["1c"]={"class_type":"LoraLoaderModelOnly","inputs":{"model":model,"lora_name":MODELS["causvid"],"strength_model":CAUSVID_STR}}; model=["1c",0] G["6"]={"class_type":"ModelSamplingSD3","inputs":{"model":model,"shift":SHIFT}}; model=["6",0] if cache and TEACACHE_NODE[0]: G["6t"]={"class_type":TEACACHE_NODE[0],"inputs":{"model":model,"model_type":"wan2.1_t2v_14B", "rel_l1_thresh":TEACACHE_THRESH,"start_percent":0.0,"end_percent":1.0,"cache_device":"cuda"}}; model=["6t",0] G["2"]={"class_type":"CLIPLoader","inputs":{"clip_name":MODELS["txt"],"type":"wan"}} G["4"]={"class_type":"CLIPTextEncode","inputs":{"clip":["2",0],"text":prompt}} G["5"]={"class_type":"CLIPTextEncode","inputs":{"clip":["2",0],"text":NEG}} G["3"]={"class_type":"VAELoader","inputs":{"vae_name":MODELS["vae"]}} G["v"]={"class_type":"VHS_LoadVideoPath","inputs":{"video":video_path,"force_rate":FPS, "frame_load_cap":FRAMES,"skip_first_frames":0,"select_every_nth":1}} G["rs"]={"class_type":"ImageScale","inputs":{"image":["v",0],"width":WIDTH,"height":HEIGHT, "upscale_method":"lanczos","crop":"center"}} G["e"]={"class_type":"VAEEncode","inputs":{"pixels":["rs",0],"vae":["3",0]}} G["8"]={"class_type":"KSampler","inputs":{"model":model,"positive":["4",0],"negative":["5",0], "latent_image":["e",0],"seed":seed,"steps":STEPS,"cfg":CFG,"sampler_name":"euler", "scheduler":"simple","denoise":max(0.1,min(1.0,strength))}} G["9"]={"class_type":"VAEDecode","inputs":{"samples":["8",0],"vae":["3",0]}} G["10"]={"class_type":"VHS_VideoCombine","inputs":{"images":["9",0],"frame_rate":FPS,"loop_count":0, "filename_prefix":prefix,"format":"video/h264-mp4","pix_fmt":"yuv420p","crf":18,"pingpong":False,"save_output":True}} return G def _spawn_v2v(prompt,video_b64,steps,strength): jid="v%d"%int(time.time()*1000) _JOBS[jid]={"status":"running","data":None,"detail":""} def _run(): try: indir=os.path.join(COMFY_DIR,"input"); os.makedirs(indir,exist_ok=True) vpath=os.path.join(indir,"v2v_%s.mp4"%jid) with open(vpath,"wb") as f: f.write(base64.b64decode(video_b64)) pr=apply_physics(prompt); stamp=int(time.time()*1000) try: path=run_graph(build_v2v(pr[:500],vpath,stamp%2**31,"v2v_%d"%stamp,strength,cache=True)) except RuntimeError as e: if any(k in str(e).lower() for k in ("rejected","comfyui error")): path=run_graph(build_v2v(pr[:500],vpath,stamp%2**31,"v2v_%d"%stamp,strength,cache=False)) else: raise with open(path,"rb") as f: data=f.read() _JOBS[jid]={"status":"done","data":data,"detail":os.path.basename(path)} try: os.remove(vpath) except Exception: pass except Exception as e: _JOBS[jid]={"status":"error","data":None,"detail":repr(e)[:200]} print(">>> v2v FAILED: %s"%repr(e)[:200],flush=True) threading.Thread(target=_run,daemon=True).start() return jid def make_app(): from fastapi import FastAPI, Request app=FastAPI() @app.get("/") def root(): return {"ok":True,"service":"lumen-lightning-worker","endpoints":["/submit","/submit_v2v","/result"]} @app.post("/submit") async def submit(req:Request): b=await req.json() return {"call_id":_spawn((b or {}).get("prompt","a cinematic scene"), int((b or {}).get("steps",STEPS)))} @app.post("/submit_v2v") async def submit_v2v(req:Request): b=await req.json() or {} vb=b.get("video_b64") if not vb: return {"error":"no video"} return {"call_id":_spawn_v2v(b.get("prompt","restyle this video"), vb, int(b.get("steps",STEPS)), float(b.get("strength",0.6)))} @app.get("/result") def result(call_id:str): j=_JOBS.get(call_id) or {"status":"unknown","data":None,"detail":"no such job"} if j["status"]=="done" and j["data"]: return {"status":"done","bytes":len(j["data"]),"video_b64":base64.b64encode(j["data"]).decode()} return {"status":j["status"],"detail":j.get("detail","")} return app if __name__=="__main__": setup(); download(); start_comfy() import uvicorn, threading # NOTEBOOK-SAFE: run uvicorn in a background thread with its OWN event loop so it # never collides with Jupyter's loop (the "asyncio loop already running" crash). # Works the same whether you paste this into a cell OR run `python lightning_worker.py`. _cfg = uvicorn.Config(make_app(), host="0.0.0.0", port=PORT, log_level="warning") _srv = uvicorn.Server(_cfg) threading.Thread(target=_srv.run, daemon=True).start() time.sleep(2) print("="*60, flush=True) print(">>> SERVING on :%d -- now expose this port in the Studio (Ports plugin)." % PORT, flush=True) print("="*60, flush=True) try: while True: time.sleep(3600) # keep alive; in a notebook the cell just stays running except KeyboardInterrupt: pass