58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
"""Native Kitchen extension: existing installation or verified first-use cache."""
|
|
import importlib
|
|
try:
|
|
from . import minimax_h3_temporal_bootstrap as bootstrap
|
|
except ImportError:
|
|
import minimax_h3_temporal_bootstrap as bootstrap
|
|
|
|
|
|
def api_for(q):
|
|
if q.device.type != 'cuda':
|
|
return None
|
|
try:
|
|
api = importlib.import_module('comfy_kitchen.temporal_attention')
|
|
except ImportError:
|
|
return bootstrap.ensure(q.device)
|
|
if getattr(api, 'VERSION', None) != 'h3_native_temporal_kernel_v1':
|
|
raise RuntimeError('Unsupported Comfy Kitchen temporal extension version; update the temporal overlay.')
|
|
return api if api.is_available(q.device) else bootstrap.ensure(q.device)
|
|
|
|
|
|
def policy_for(api, plan, ctx):
|
|
key=(api.VERSION,ctx.strength,ctx.window,ctx.falloff,ctx.cap)
|
|
if getattr(plan,'_native_temporal_policy_key',None)!=key:
|
|
times = getattr(plan, 'frame_times', ())
|
|
if not times:
|
|
times = tuple(ctx.times[0, plan.target_start:plan.target_stop:plan.target_frame_rows].detach().cpu().tolist())
|
|
plan._native_temporal_policy = api.TemporalPolicy(
|
|
plan.source_start, plan.source_stop, plan.target_start, plan.target_stop,
|
|
plan.source_frame_rows, plan.target_frame_rows, tuple(times),
|
|
ctx.strength, ctx.window, ctx.falloff, ctx.cap)
|
|
plan._native_temporal_policy_key=key
|
|
return plan._native_temporal_policy
|
|
|
|
|
|
def run(kind, q, k, v, plan, ctx, **kwargs):
|
|
api = api_for(q)
|
|
if api is None:
|
|
return None
|
|
policy = policy_for(api, plan, ctx)
|
|
if not getattr(plan, '_native_'+kind+'_logged', False):
|
|
print(f'[BFS H3 native temporal] Kitchen ONE-PASS {kind} kernel ACTIVE; temporal bias fused; no per-frame attention loop.', flush=True)
|
|
setattr(plan, '_native_'+kind+'_logged', True)
|
|
dtype=q.dtype
|
|
if str(dtype)=='torch.float32':
|
|
import torch
|
|
q,k,v=[t.to(torch.bfloat16) for t in (q,k,v)]
|
|
fn = api.int8_attention if kind == 'dense' else api.sol_attn
|
|
return fn(q,k,v,policy,**kwargs).to(dtype)
|
|
|
|
|
|
def prepare_model(model):
|
|
"""Prepare a selected native backend during node execution, before sampling."""
|
|
device=getattr(model,'load_device',None)
|
|
if getattr(device,'type',None)=='cuda':
|
|
from types import SimpleNamespace
|
|
return api_for(SimpleNamespace(device=device))
|
|
return None
|