From 032bf6972af4ba598a9e3a3ffd3925d0789cdc97 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 29 Aug 2025 10:34:40 -0400 Subject: [PATCH] uniform hf auth check Signed-off-by: Vladimir Mandic --- modules/sd_models.py | 2 +- pipelines/model_hdm.py | 1 + pipelines/model_kolors.py | 8 +++--- pipelines/model_lumina.py | 6 +++-- pipelines/model_meissonic.py | 13 +++++---- pipelines/model_omnigen.py | 51 +++++++++++++++++++++++++++++++++--- pipelines/model_omnigen2.py | 48 --------------------------------- pipelines/model_pixart.py | 2 +- pipelines/model_sana.py | 2 +- 9 files changed, 67 insertions(+), 66 deletions(-) delete mode 100644 pipelines/model_omnigen2.py diff --git a/modules/sd_models.py b/modules/sd_models.py index ffaec5656..2f12a1684 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -339,7 +339,7 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op=' sd_model = load_meissonic(checkpoint_info, diffusers_load_config) allow_post_quant = True elif model_type in ['OmniGen2']: # forced pipeline - from pipelines.model_omnigen2 import load_omnigen2 + from pipelines.model_omnigen import load_omnigen2 sd_model = load_omnigen2(checkpoint_info, diffusers_load_config) allow_post_quant = False elif model_type in ['OmniGen']: # forced pipeline diff --git a/pipelines/model_hdm.py b/pipelines/model_hdm.py index 823a2b37f..6b6348dfb 100644 --- a/pipelines/model_hdm.py +++ b/pipelines/model_hdm.py @@ -6,6 +6,7 @@ from modules import shared, devices, sd_models, errors def load_hdm(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) try: devices.dtype = torch.float16 diff --git a/pipelines/model_kolors.py b/pipelines/model_kolors.py index c522d1605..fe95e2b4a 100644 --- a/pipelines/model_kolors.py +++ b/pipelines/model_kolors.py @@ -1,14 +1,16 @@ import torch import diffusers -from modules import shared, devices, sd_hijack_te +from modules import shared, devices, sd_models, sd_hijack_te -def load_kolors(_checkpoint_info, diffusers_load_config={}): +def load_kolors(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + diffusers_load_config['variant'] = "fp16" if 'torch_dtype' not in diffusers_load_config: diffusers_load_config['torch_dtype'] = torch.float16 - repo_id = 'Kwai-Kolors/Kolors-diffusers' shared.log.debug(f'Load model: type=Kolors repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') pipe = diffusers.KolorsPipeline.from_pretrained( repo_id, diff --git a/pipelines/model_lumina.py b/pipelines/model_lumina.py index a75b999a0..060cd3f0b 100644 --- a/pipelines/model_lumina.py +++ b/pipelines/model_lumina.py @@ -1,12 +1,13 @@ import transformers import diffusers -from modules import shared, sd_models, sd_hijack_te, devices, modelloader, model_quant +from modules import shared, sd_models, sd_hijack_te, devices, model_quant from pipelines import generic def load_lumina(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) - modelloader.hf_login() + sd_models.hf_auth_check(checkpoint_info) + load_config, _quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) shared.log.debug(f'Load model: type=LuminaSFT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') pipe = diffusers.LuminaText2ImgPipeline.from_pretrained( @@ -21,6 +22,7 @@ def load_lumina(checkpoint_info, diffusers_load_config={}): def load_lumina2(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) if shared.opts.teacache_enabled: from modules import teacache diff --git a/pipelines/model_meissonic.py b/pipelines/model_meissonic.py index e5654e85f..06bb41b07 100644 --- a/pipelines/model_meissonic.py +++ b/pipelines/model_meissonic.py @@ -11,9 +11,8 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}): from pipelines.meissonic.pipeline_inpaint import MeissonicInpaintPipeline shared_items.pipelines['Meissonic'] = MeissonicPipeline - modelloader.hf_login() repo_id = sd_models.path_to_repo(checkpoint_info) - cache_dir = shared.opts.diffusers_dir + sd_models.hf_auth_check(checkpoint_info) diffusers_load_config['variant'] = 'fp16' diffusers_load_config['trust_remote_code'] = True @@ -22,29 +21,29 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}): model = TransformerMeissonic.from_pretrained( repo_id, subfolder="transformer", - cache_dir=cache_dir, + cache_dir=shared.opts.diffusers_dir, **diffusers_load_config, ) vqvae = diffusers.VQModel.from_pretrained( repo_id, subfolder="vqvae", - cache_dir=cache_dir, + cache_dir=shared.opts.diffusers_dir, **diffusers_load_config, ) text_encoder = transformers.CLIPTextModelWithProjection.from_pretrained( repo_id, subfolder="text_encoder", - cache_dir=cache_dir, + cache_dir=shared.opts.diffusers_dir, ) tokenizer = transformers.CLIPTokenizer.from_pretrained( repo_id, subfolder="tokenizer", - cache_dir=cache_dir, + cache_dir=shared.opts.diffusers_dir, ) scheduler = MeissonicScheduler.from_pretrained( repo_id, subfolder="scheduler", - cache_dir=cache_dir, + cache_dir=shared.opts.diffusers_dir, ) pipe = MeissonicPipeline( vqvae=vqvae.to(devices.dtype), diff --git a/pipelines/model_omnigen.py b/pipelines/model_omnigen.py index bdc091341..6cff7de6f 100644 --- a/pipelines/model_omnigen.py +++ b/pipelines/model_omnigen.py @@ -4,7 +4,7 @@ from modules import shared, devices, sd_models, model_quant, sd_hijack_te def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument repo_id = sd_models.path_to_repo(checkpoint_info) - vae = None + sd_models.hf_auth_check(checkpoint_info) load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model') shared.log.debug(f'Load model: type=OmniGen repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') @@ -17,8 +17,6 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u ) load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) - if vae is not None: - load_config['vae'] = vae pipe = diffusers.OmniGenPipeline.from_pretrained( repo_id, transformer=transformer, @@ -30,3 +28,50 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u devices.torch_gc(force=True, reason='load') return pipe + + +def load_omnigen2(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + from pipelines.omnigen2 import OmniGen2Pipeline, OmniGen2Transformer2DModel, Qwen2_5_VLForConditionalGeneration + diffusers.OmniGen2Pipeline = OmniGen2Pipeline # monkey-pathch + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline + + load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model') + shared.log.debug(f'Load model: type=OmniGen2 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') + transformer = OmniGen2Transformer2DModel.from_pretrained( + repo_id, + subfolder="transformer", + cache_dir=shared.opts.diffusers_dir, + trust_remote_code=True, + **load_config, + **quant_config, + ) + + load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='TE') + mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained( + repo_id, + subfolder="mllm", + cache_dir=shared.opts.diffusers_dir, + trust_remote_code=True, + **load_config, + **quant_config, + ) + + pipe = OmniGen2Pipeline.from_pretrained( + repo_id, + # transformer=transformer, + mllm=mllm, + cache_dir=shared.opts.diffusers_dir, + trust_remote_code=True, + **load_config, + ) + pipe.transformer = transformer # for omnigen2 transformer must be loaded after pipeline + + sd_hijack_te.init_hijack(pipe) + + devices.torch_gc(force=True, reason='load') + return pipe diff --git a/pipelines/model_omnigen2.py b/pipelines/model_omnigen2.py deleted file mode 100644 index 5a4fc5b6e..000000000 --- a/pipelines/model_omnigen2.py +++ /dev/null @@ -1,48 +0,0 @@ -import diffusers -from modules import shared, devices, sd_models, model_quant, sd_hijack_te - - -def load_omnigen2(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument - repo_id = sd_models.path_to_repo(checkpoint_info) - - from pipelines.omnigen2 import OmniGen2Pipeline, OmniGen2Transformer2DModel, Qwen2_5_VLForConditionalGeneration - diffusers.OmniGen2Pipeline = OmniGen2Pipeline # monkey-pathch - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline - - load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=OmniGen2 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') - transformer = OmniGen2Transformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.diffusers_dir, - trust_remote_code=True, - **load_config, - **quant_config, - ) - - load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='TE') - mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained( - repo_id, - subfolder="mllm", - cache_dir=shared.opts.diffusers_dir, - trust_remote_code=True, - **load_config, - **quant_config, - ) - - pipe = OmniGen2Pipeline.from_pretrained( - repo_id, - # transformer=transformer, - mllm=mllm, - cache_dir=shared.opts.diffusers_dir, - trust_remote_code=True, - **load_config, - ) - pipe.transformer = transformer # for omnigen2 transformer must be loaded after pipeline - - sd_hijack_te.init_hijack(pipe) - - devices.torch_gc(force=True, reason='load') - return pipe diff --git a/pipelines/model_pixart.py b/pipelines/model_pixart.py index e0f5f7409..b628e34ec 100644 --- a/pipelines/model_pixart.py +++ b/pipelines/model_pixart.py @@ -33,7 +33,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}): del text_encoder del transformer - # sd_hijack_te.init_hijack(pipe) + sd_hijack_te.init_hijack(pipe) devices.torch_gc(force=True, reason='load') return pipe diff --git a/pipelines/model_sana.py b/pipelines/model_sana.py index 4abd17620..ce101cb37 100644 --- a/pipelines/model_sana.py +++ b/pipelines/model_sana.py @@ -22,8 +22,8 @@ def load_quants(kwargs, repo_id, cache_dir): def load_sana(checkpoint_info, kwargs={}): - modelloader.hf_login() repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) kwargs.pop('load_connected_pipeline', None) kwargs.pop('safety_checker', None)