From 60c9b04481dcfc7ced07d80f50b1ac6b14d4e4bb Mon Sep 17 00:00:00 2001 From: Stan Travis <160105964+stantravis@users.noreply.github.com> Date: Fri, 3 Jan 2025 17:36:06 +0300 Subject: [PATCH 1/4] Added fp8 mode for ComfyUI, update comfyui_nodes.py --- comfyui/comfyui_nodes.py | 50 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 46 insertions(+), 4 deletions(-) diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index 3af9285..dd5c2b4 100644 --- a/comfyui/comfyui_nodes.py +++ b/comfyui/comfyui_nodes.py @@ -51,7 +51,16 @@ class Ruyi_LoadModel: { "default": "yes" } - ) + ), + + "fp8_quant_mode": ( + [ "none", "lite", "strong", "extreme"], + { "default": "none", "tooltip": 'Mode "extreme" is not recommended for good quality'} + ), + "fp8_data_type": ( + [ "auto", "fp8_e4m3fn", "fp8_e5m2"], + { "default": "auto"} + ), }, } @@ -70,7 +79,7 @@ class Ruyi_LoadModel: def try_setup_pipeline(self, model_path, weight_dtype, config): try: # Init processbar - pbar = ProgressBar(5) + pbar = ProgressBar(6) # Get Vae vae = AutoencoderKLMagvit.from_pretrained( @@ -90,6 +99,31 @@ class Ruyi_LoadModel: # Update pbar pbar.update(1) + if self.fp8_quant_mode != 'none': + count_f8 =0 + fp8_type = torch.float8_e5m2 if self.fp8_data_type == 'fp8_e5m2' else torch.float8_e4m3fn + if self.fp8_quant_mode != 'extreme': + + shape_size = 2816 ** 2 + if self.fp8_quant_mode == 'strong': shape_size -= 1 + + for module in transformer.modules(): + if module.__class__.__name__ in ["Linear"]: + x,y = module.weight.shape + if x * y > shape_size: + module.to(fp8_type) + count_f8 += 1 + else: + for module in transformer.modules(): + if len(list(module.modules())) > 1: continue + elif module.__class__.__name__ not in ["Embedding"]: + module.to(fp8_type) + count_f8 += 1 + + print (f'FP8: {count_f8} layers converted to {fp8_type}') + # Update pbar + pbar.update(1) + # Load Clip clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained( model_path, subfolder="image_encoder" @@ -126,8 +160,11 @@ class Ruyi_LoadModel: except Exception as e: print("[Ruyi] Setup pipeline failed:", e) return None + + def load_model(self, model, auto_download, auto_update, fp8_quant_mode='none', fp8_data_type='auto'): + self.fp8_quant_mode = fp8_quant_mode + self.fp8_data_type = fp8_data_type - def load_model(self, model, auto_download, auto_update): # Init weight_dtype and device device = mm.get_torch_device() offload_device = mm.unet_offload_device() @@ -387,8 +424,13 @@ class Ruyi_I2VSampler: # Load control embeddings embeddings = self.get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction) + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + # Inference - with torch.no_grad(): + with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype): video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width)) From 963f2c9bf4e4c69811003438530067fc467607d7 Mon Sep 17 00:00:00 2001 From: Stan Travis <160105964+stantravis@users.noreply.github.com> Date: Fri, 3 Jan 2025 17:44:48 +0300 Subject: [PATCH 2/4] Add fp8 info, Update README.md --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index bf6fb01..0e4bc00 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,8 @@ +**New**: +FP8 added for ComfyUI node. + GPU memory decreases depending on the modes: bf16 default > fp8 lite > fp8 strong > fp8 extreme + +----------------- # Ruyi-Models English | [简体中文](./README_CN.md) From e2a64fbe35dcd0b2f3eeb9e38e2c3e4b5a7d2f35 Mon Sep 17 00:00:00 2001 From: Stan Travis <160105964+stantravis@users.noreply.github.com> Date: Fri, 3 Jan 2025 20:21:38 +0300 Subject: [PATCH 3/4] Deleting meaningless --- comfyui/comfyui_nodes.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index dd5c2b4..646802e 100644 --- a/comfyui/comfyui_nodes.py +++ b/comfyui/comfyui_nodes.py @@ -424,11 +424,6 @@ class Ruyi_I2VSampler: # Load control embeddings embeddings = self.get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction) - try: - torch.cuda.reset_peak_memory_stats(device) - except: - pass - # Inference with torch.no_grad(), torch.autocast(str(device), dtype = pipeline.transformer.dtype): video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 From 7dd5e9d871fb352ccfae736175a95d3997068982 Mon Sep 17 00:00:00 2001 From: Stan Travis <160105964+stantravis@users.noreply.github.com> Date: Sun, 5 Jan 2025 04:18:05 +0300 Subject: [PATCH 4/4] Better quality for mode "extreme" (with fp8_e4m3fn). Update comfyui_nodes.py For "extreme" mode, the code has been corrected to be more accurate and layers affecting quality have been exclude. --- comfyui/comfyui_nodes.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index 646802e..7159cc0 100644 --- a/comfyui/comfyui_nodes.py +++ b/comfyui/comfyui_nodes.py @@ -55,11 +55,11 @@ class Ruyi_LoadModel: "fp8_quant_mode": ( [ "none", "lite", "strong", "extreme"], - { "default": "none", "tooltip": 'Mode "extreme" is not recommended for good quality'} + { "default": "none"} ), "fp8_data_type": ( [ "auto", "fp8_e4m3fn", "fp8_e5m2"], - { "default": "auto"} + { "default": "auto", "tooltip": 'do not use "fp8_e5m2" with "extreme" mode'} ), }, } @@ -100,7 +100,7 @@ class Ruyi_LoadModel: pbar.update(1) if self.fp8_quant_mode != 'none': - count_f8 =0 + count_f8 = 0 fp8_type = torch.float8_e5m2 if self.fp8_data_type == 'fp8_e5m2' else torch.float8_e4m3fn if self.fp8_quant_mode != 'extreme': @@ -115,11 +115,11 @@ class Ruyi_LoadModel: count_f8 += 1 else: for module in transformer.modules(): - if len(list(module.modules())) > 1: continue - elif module.__class__.__name__ not in ["Embedding"]: - module.to(fp8_type) - count_f8 += 1 - + if len(list(module.modules())) == 1 and list(module.named_parameters()): + if module.__class__.__name__ not in ["Embedding",'LayerNorm','Conv2d','NonDynamicallyQuantizableLinear']: + module.to(fp8_type) + count_f8 += 1 + print (f'FP8: {count_f8} layers converted to {fp8_type}') # Update pbar pbar.update(1)