Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

add experimental pruna

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2026-07-02 11:25:33 +02:00
parent bc86a611de
commit fcee7e23f2
7 changed files with 164 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/postprocess/nvvfx_model.py
+104
View File
@@ -0,0 +1,104 @@
import torch
import numpy as np
from PIL import Image
from modules import shared, devices
from modules.logger import log
from modules.upscaler import Upscaler, UpscalerData
class UpscalerDiffusion(Upscaler):
def __init__(self, dirname): # pylint: disable=super-init-not-called
self.name = "nVidia VFX"
self.user_path = dirname
"""
self.scalers = [
UpscalerData(name="nVidia VFX 1x Denoise Ultra", path="", upscaler=self, model=None, scale=1),
UpscalerData(name="nVidia VFX 1x Deblur Ultra", path="", upscaler=self, model=None, scale=1),
UpscalerData(name="nVidia VFX 1x Denoise High", path="", upscaler=self, model=None, scale=1),
UpscalerData(name="nVidia VFX 1x Deblur High", path="", upscaler=self, model=None, scale=1),
UpscalerData(name="nVidia VFX 2x Ultra", path="", upscaler=self, model=None, scale=2),
UpscalerData(name="nVidia VFX 4x Ultra", path="", upscaler=self, model=None, scale=4),
UpscalerData(name="nVidia VFX 2x High", path="", upscaler=self, model=None, scale=2),
UpscalerData(name="nVidia VFX 4x High", path="", upscaler=self, model=None, scale=4),
]
"""
self.scalers = []
self.models = {}
def load_model(self, path: str):
scaler: UpscalerData = [x for x in self.scalers if x.data_path == path or x.name == path]
if len(scaler) == 0:
log.error(f"Upscaler cannot match model: type={self.name} model={path}")
return None
scaler = scaler[0]
if self.models.get(path, None) is not None:
log.debug(f"Upscaler cached: type={scaler.name} model={path}")
return self.models[path]
from installer import install
install('nvidia-vfx')
def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor):
pass
def do_upscale(self, img: Image.Image, selected_model):
devices.torch_gc()
self.load_model(selected_model)
frame = torch.from_numpy(np.array(img)).permute(2, 0, 1).float().to(devices.device) / 255.0
frame = frame.to(devices.device)
try:
import nvvfx
except Exception as e:
log.error(f"Upscaler: failed to import nvvfx: {e}")
return img
config_map = {
"nVidia VFX 1x Denoise Ultra": nvvfx.VideoSuperRes.QualityLevel.DENOISE_ULTRA,
"nVidia VFX 1x Deblur Ultra": nvvfx.VideoSuperRes.QualityLevel.DEBLUR_ULTRA,
"nVidia VFX 1x Denoise High": nvvfx.VideoSuperRes.QualityLevel.DENOISE_HIGH,
"nVidia VFX 1x Deblur High": nvvfx.VideoSuperRes.QualityLevel.DEBLUR_HIGH,
"nVidia VFX 2x Ultra": nvvfx.VideoSuperRes.QualityLevel.ULTRA,
"nVidia VFX 4x Ultra": nvvfx.VideoSuperRes.QualityLevel.ULTRA,
"nVidia VFX 2x High": nvvfx.VideoSuperRes.QualityLevel.HIGH,
"nVidia VFX 4x High": nvvfx.VideoSuperRes.QualityLevel.HIGH,
}
quality = config_map.get(selected_model, None)
log.info(f'Upscaler: type="{self.name}" model="{selected_model}" version={nvvfx.__version__} sdk={nvvfx.get_sdk_version()} quality={quality}')
if self.models.get(selected_model, None) is not None:
vsr = self.models[selected_model]
else:
vsr = nvvfx.VideoSuperRes(quality=quality)
self.models[selected_model] = vsr
if '2x' in selected_model:
vsr.output_width = img.width * 2
vsr.output_height = img.height * 2
elif '4x' in selected_model:
vsr.output_width = img.width * 4
vsr.output_height = img.height * 4
elif 'Denoise' in selected_model or 'Deblur' in selected_model or '1x' in selected_model:
vsr.output_width = img.width
vsr.output_height = img.height
else:
log.error(f"Upscaler: unknown model: {selected_model}")
return img
vsr.input_width = img.width
vsr.input_height = img.height
log.debug(f"Upscaler: {vsr}")
try:
vsr.load()
except Exception as e:
log.error(f"Upscaler: failed to load model: {selected_model} error={e}")
return img
self.models[selected_model] = vsr
result = vsr.run(frame)
result = torch.from_dlpack(result.image).clone()
image = Image.fromarray((result.permute(1, 2, 0).contiguous().cpu().numpy() * 255).astype(np.uint8))
if shared.opts.upscaler_unload and selected_model in self.models:
del self.models[selected_model]
log.debug(f"Upscaler unloaded: type={self.name} model={selected_model}")
devices.torch_gc(force=True)
return image
modules/postprocess/realesrgan_model.py
+1
View File
@@ -6,6 +6,7 @@ from modules.upscaler import Upscaler
from modules.shared import opts, device, log
from modules import devices
class UpscalerRealESRGAN(Upscaler):
def __init__(self, dirname):
from installer import install
modules/sd_models_compile.py
+37 -6
View File
@@ -87,14 +87,45 @@ def optimize_openvino(sd_model, clear_cache=True):
def compile_pruna(sd_model):
# TODO pruna: enable when it supports transformers==5.5
# install('pruna')
"""
import warnings
from installer import install
install('pruna')
# pip install pruna[stable-fast] --extra-index-url https://prunaai.pythonanywhere.com/
from pruna import smash, SmashConfig
smash_config = SmashConfig(["deepcache", "stable_fast"])
smashed_model = smash(model=sd_model, smash_config=smash_config)
return smashed_model
# https://docs.pruna.ai/en/stable/compression.html
"""
cachers = ["fastercache", "deepcache", "fora", "pab"]
compilers = ["stable_fast", "x_fast", "torch_compile"]
factorizers = ["qkv_diffusers"]
pruners = ["kvpress", "padding_pruning", "token_merging", "torch_structured", "torch_unstructured"]
kernels = ["flash_attn3", "ring_attn", "sage_attn"]
distillers = ["text_to_image_distillation_inplace_perp", "text_to_image_distillation_lora", "text_to_image_distillation_perp", "hyper"]
enhancers = ["img2img_denoise", "realesrgan_upscale"]
quants = ["c_generate", "c_translate", "c_whisper", "llama_cpp"]
quantizers = ["gptq", "half", "hqq", "hqq_diffusers", "diffusers_int8", "awq", "torch_dynamic", "torchao"]
"""
config_list = shared.opts.pruna_cachers + shared.opts.pruna_compilers + shared.opts.pruna_factorizers + shared.opts.pruna_pruners
if len(config_list) == 0:
log.warning(f"Model compile: task=pruna pipeline={sd_model.__class__.__name__} no algorithms selected")
return sd_model
config = SmashConfig(configuration=config_list, device=devices.device)
log.info(f"Model compile: task=pruna pipeline={sd_model.__class__.__name__} config={config}")
try:
smashed_model = smash(
model=sd_model,
smash_config=config,
experimental=shared.opts.pruna_experimental,
)
return smashed_model
except Exception as e:
log.error(f"Model compile: task=pruna pipeline={sd_model.__class__.__name__} error={e}")
errors.display(e, 'Compile')
finally:
# re-silence warnings after pruna compile, as it enables a lot of warnings
warnings.filterwarnings(action="ignore", category=DeprecationWarning)
warnings.filterwarnings(action="ignore", category=FutureWarning)
warnings.filterwarnings(action="ignore", category=UserWarning)
return sd_model
modules/ui_definitions.py
+9 -1
View File
@@ -414,9 +414,17 @@ def create_settings(cmd_opts):
"cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
"cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler", "Control"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE", "LLM", "Control", "Upscaler"]}),
"cuda_compile_backend": OptionInfo("inductor" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'migraphx', 'ipex', 'onediff', 'stable-fast', 'deep-cache', 'olive-ai', 'openvino', 'openvino_fx', 'pruna']}),
"torch_compile_sep": OptionInfo("<h2>Torch Compile</h2>", "", gr.HTML),
"cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}),
"cuda_compile_options": OptionInfo(["repeated", "dynamic", "components"] if not cmd_opts.use_openvino else [], "Model compile options", gr.CheckboxGroup, {"choices": ["components", "precompile", "repeated", "fullgraph", "dynamic", "verbose"]}),
"cuda_compile_options": OptionInfo(["repeated", "dynamic", "components"] if not cmd_opts.use_openvino else [], "Torch compile options", gr.CheckboxGroup, {"choices": ["components", "precompile", "repeated", "fullgraph", "dynamic", "verbose"]}),
"deepcache_compile_sep": OptionInfo("<h2>DeepCache</h2>", "", gr.HTML),
"deep_cache_interval": OptionInfo(3, "DeepCache cache interval", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
"pruna_compile_sep": OptionInfo("<h2>Pruna</h2>", "", gr.HTML),
"pruna_experimental": OptionInfo(False, "Pruna experimental features", gr.Checkbox),
"pruna_cachers": OptionInfo([], "Pruna cachers", gr.CheckboxGroup, {"choices": ["fastercache", "deepcache", "fora", "pab"]}),
"pruna_compilers": OptionInfo([], "Pruna compilers", gr.CheckboxGroup, {"choices": ["stable_fast", "x_fast", "torch_compile"]}),
"pruna_factorizers": OptionInfo([], "Pruna factorizers", gr.CheckboxGroup, {"choices": ["qkv_diffusers"]}),
"pruna_pruners": OptionInfo([], "Pruna pruners", gr.CheckboxGroup, {"choices": ["kvpress", "padding_pruning", "token_merging", "torch_structured", "torch_unstructured"]}),
}))
# --- System Paths ---