From d3292f2f059bf458494313ca4127908d5f4e2025 Mon Sep 17 00:00:00 2001 From: smthemex <138738845+smthemex@users.noreply.github.com> Date: Fri, 29 Nov 2024 21:21:36 +0800 Subject: [PATCH] Update app_flux.py --- PuLID/app_flux.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/PuLID/app_flux.py b/PuLID/app_flux.py index b158de1..d8f4448 100644 --- a/PuLID/app_flux.py +++ b/PuLID/app_flux.py @@ -8,7 +8,7 @@ from .flux.util import load_ae, load_clip, load_flow_model, load_t5, load_flow_m from .pulid.pipeline_flux import PuLIDPipeline -def get_models(name: str,ckpt_path,vae_cf, clip_cf,if_repo,aggressive_offload,device: torch.device, offload: bool,quantized_mode): +def get_models(name: str,ckpt_path,vae_cf, clip_cf,if_repo,use_quantize,device: torch.device, offload: bool,quantized_mode): t5 = load_t5(name,clip_cf,if_repo,device, max_length=128) clip = load_clip(name,clip_cf,if_repo,device) if if_repo: @@ -18,7 +18,7 @@ def get_models(name: str,ckpt_path,vae_cf, clip_cf,if_repo,aggressive_offload,de else: name = "flux-schnell" if quantized_mode=="fp8" or quantized_mode=="nf4": - model = load_flow_model_quintized(name, ckpt_path,quantized_mode,aggressive_offload,device="cpu" if offload else device) + model = load_flow_model_quintized(name, ckpt_path,quantized_mode,use_quantize,device="cpu" if offload else device) else: model = load_flow_model(name, ckpt_path,device="cpu" if offload else device) model.eval() @@ -27,7 +27,7 @@ def get_models(name: str,ckpt_path,vae_cf, clip_cf,if_repo,aggressive_offload,de class FluxGenerator: - def __init__(self, model_name: str, ckpt_path,device: str, offload: bool,aggressive_offload: bool, pretrained_model,quantized_mode,clip_vision_path,clip_cf,vae_cf,if_repo,onnx_provider): + def __init__(self, model_name: str, ckpt_path,device: str, offload: bool,aggressive_offload: bool, pretrained_model,quantized_mode,clip_vision_path,clip_cf,vae_cf,if_repo,onnx_provider,use_quantize): self.device = torch.device(device) self.offload = offload self.model_name = model_name @@ -35,12 +35,13 @@ class FluxGenerator: self.ckpt_path=ckpt_path self.clip_vision_path=clip_vision_path self.clip_cf=clip_cf + self.use_quantize=use_quantize self.vae_cf=vae_cf self.if_repo=if_repo self.onnx_provider=onnx_provider self.quantized_mode=quantized_mode self.model, self.ae, self.t5, self.clip = get_models( - model_name,self.ckpt_path,self.vae_cf,self.clip_cf,self.if_repo,self.aggressive_offload, + model_name,self.ckpt_path,self.vae_cf,self.clip_cf,self.if_repo,self.use_quantize, device=self.device, offload=self.offload,quantized_mode=self.quantized_mode, )