Ethanfel
0fecfcee37
Vendor minimal diffsynth subset for FlashVSR inference (full/tiny pipelines, v1 and v1.1 checkpoints auto-downloaded from HuggingFace). Includes segment-based processing with temporal overlap and crossfade blending for bounded RAM on long videos. Nodes: Load FlashVSR Model, FlashVSR Upscale, FlashVSR Segment Upscale. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
96 lines
4.7 KiB
Python
96 lines
4.7 KiB
Python
import torch, copy
|
|
from ..models.utils import init_weights_on_device
|
|
|
|
|
|
def cast_to(weight, dtype, device):
|
|
r = torch.empty_like(weight, dtype=dtype, device=device)
|
|
r.copy_(weight)
|
|
return r
|
|
|
|
|
|
class AutoWrappedModule(torch.nn.Module):
|
|
def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device):
|
|
super().__init__()
|
|
self.module = module.to(dtype=offload_dtype, device=offload_device)
|
|
self.offload_dtype = offload_dtype
|
|
self.offload_device = offload_device
|
|
self.onload_dtype = onload_dtype
|
|
self.onload_device = onload_device
|
|
self.computation_dtype = computation_dtype
|
|
self.computation_device = computation_device
|
|
self.state = 0
|
|
|
|
def offload(self):
|
|
if self.state == 1 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device):
|
|
self.module.to(dtype=self.offload_dtype, device=self.offload_device)
|
|
self.state = 0
|
|
|
|
def onload(self):
|
|
if self.state == 0 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device):
|
|
self.module.to(dtype=self.onload_dtype, device=self.onload_device)
|
|
self.state = 1
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device:
|
|
module = self.module
|
|
else:
|
|
module = copy.deepcopy(self.module).to(dtype=self.computation_dtype, device=self.computation_device)
|
|
return module(*args, **kwargs)
|
|
|
|
|
|
class AutoWrappedLinear(torch.nn.Linear):
|
|
def __init__(self, module: torch.nn.Linear, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device):
|
|
with init_weights_on_device(device=torch.device("meta")):
|
|
super().__init__(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None, dtype=offload_dtype, device=offload_device)
|
|
self.weight = module.weight
|
|
self.bias = module.bias
|
|
self.offload_dtype = offload_dtype
|
|
self.offload_device = offload_device
|
|
self.onload_dtype = onload_dtype
|
|
self.onload_device = onload_device
|
|
self.computation_dtype = computation_dtype
|
|
self.computation_device = computation_device
|
|
self.state = 0
|
|
|
|
def offload(self):
|
|
if self.state == 1 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device):
|
|
self.to(dtype=self.offload_dtype, device=self.offload_device)
|
|
self.state = 0
|
|
|
|
def onload(self):
|
|
if self.state == 0 and (self.offload_dtype != self.onload_dtype or self.offload_device != self.onload_device):
|
|
self.to(dtype=self.onload_dtype, device=self.onload_device)
|
|
self.state = 1
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device:
|
|
weight, bias = self.weight, self.bias
|
|
else:
|
|
weight = cast_to(self.weight, self.computation_dtype, self.computation_device)
|
|
bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device)
|
|
return torch.nn.functional.linear(x, weight, bias)
|
|
|
|
|
|
def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0):
|
|
for name, module in model.named_children():
|
|
for source_module, target_module in module_map.items():
|
|
if isinstance(module, source_module):
|
|
num_param = sum(p.numel() for p in module.parameters())
|
|
if max_num_param is not None and total_num_param + num_param > max_num_param:
|
|
module_config_ = overflow_module_config
|
|
else:
|
|
module_config_ = module_config
|
|
module_ = target_module(module, **module_config_)
|
|
setattr(model, name, module_)
|
|
total_num_param += num_param
|
|
break
|
|
else:
|
|
total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param)
|
|
return total_num_param
|
|
|
|
|
|
def enable_vram_management(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None):
|
|
enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0)
|
|
model.vram_management_enabled = True
|
|
|