From c31b19cb577a3dcab1ecf4ab814b993ffb5b47f2 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 13 Jun 2024 19:10:29 -0400 Subject: [PATCH] sd3 simplify safetensors handler --- modules/model_sd3.py | 141 ++++++++++++++----------------------------- 1 file changed, 44 insertions(+), 97 deletions(-) diff --git a/modules/model_sd3.py b/modules/model_sd3.py index 0b7373572..614c37f6a 100644 --- a/modules/model_sd3.py +++ b/modules/model_sd3.py @@ -27,7 +27,7 @@ def hf_login(): loggedin = True -def load_sd3(te3=None, fn=None, cache_dir=None, config=None): +def load_sd3(fn=None, cache_dir=None, config=None): from modules import devices hf_login() repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers' @@ -37,78 +37,53 @@ def load_sd3(te3=None, fn=None, cache_dir=None, config=None): if fn is not None and fn.endswith('.safetensors') and os.path.exists(fn): model_id = fn loader = diffusers.StableDiffusion3Pipeline.from_single_file - kwargs = { - 'text_encoder': transformers.CLIPTextModelWithProjection.from_pretrained( - repo_id, - subfolder='text_encoder', - cache_dir=cache_dir, - torch_dtype=dtype, - ), - 'text_encoder_2': transformers.CLIPTextModelWithProjection.from_pretrained( - repo_id, - subfolder='text_encoder_2', - cache_dir=cache_dir, - torch_dtype=dtype, - ), - 'tokenizer': transformers.CLIPTokenizer.from_pretrained( - repo_id, - subfolder='tokenizer', - cache_dir=cache_dir, - ), - 'tokenizer_2': transformers.CLIPTokenizer.from_pretrained( - repo_id, - subfolder='tokenizer_2', - cache_dir=cache_dir, - ), - } + diffusers_minor = int(diffusers.__version__.split('.')[1]) + fn_size = os.path.getsize(fn) + if diffusers_minor < 30 or fn_size < 5e9: # te1/te2 do not get loaded correctly in diffusers 0.29.0 or model is without te1/te2 + kwargs = { + 'text_encoder': transformers.CLIPTextModelWithProjection.from_pretrained( + repo_id, + subfolder='text_encoder', + cache_dir=cache_dir, + torch_dtype=dtype, + ), + 'text_encoder_2': transformers.CLIPTextModelWithProjection.from_pretrained( + repo_id, + subfolder='text_encoder_2', + cache_dir=cache_dir, + torch_dtype=dtype, + ), + 'tokenizer': transformers.CLIPTokenizer.from_pretrained( + repo_id, + subfolder='tokenizer', + cache_dir=cache_dir, + ), + 'tokenizer_2': transformers.CLIPTokenizer.from_pretrained( + repo_id, + subfolder='tokenizer_2', + cache_dir=cache_dir, + ), + 'text_encoder_3': None, + } + elif fn_size < 1e10: # if model is below 10gb it does not have te4 + kwargs = { + 'text_encoder_3': None, + } + else: + kwargs = {} else: model_id = repo_id loader = diffusers.StableDiffusion3Pipeline.from_pretrained - if te3 == 'fp16': - text_encoder_3 = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder='text_encoder_3', - torch_dtype=dtype, - cache_dir=cache_dir, - ) - pipe = loader( - model_id, - torch_dtype=dtype, - text_encoder_3=text_encoder_3, - cache_dir=cache_dir, - config=config, - **kwargs, - ) - elif te3 == 'fp8': - quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True) - text_encoder_3 = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder='text_encoder_3', - quantization_config=quantization_config, - cache_dir=cache_dir, - config=config, - ) - pipe = loader( - model_id, - text_encoder_3=text_encoder_3, - device_map='balanced', - torch_dtype=dtype, - cache_dir=cache_dir, - config=config, - **kwargs, - ) - else: - pipe = loader( - model_id, - torch_dtype=dtype, - text_encoder_3=None, - cache_dir=cache_dir, - config=config, - **kwargs, - ) + pipe = loader( + model_id, + torch_dtype=dtype, + cache_dir=cache_dir, + config=config, + **kwargs, + ) diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["stable-diffusion-3"] = diffusers.StableDiffusion3Pipeline diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["stable-diffusion-3"] = diffusers.StableDiffusion3Img2ImgPipeline - devices.torch_gc(force=True) + devices.torch_gc() return pipe @@ -145,32 +120,4 @@ def load_te3(pipe, te3=None, cache_dir=None): subfolder='tokenizer_3', cache_dir=cache_dir, ) - devices.torch_gc(force=True) - - -if __name__ == '__main__': - model_fn = '/mnt/models/stable-diffusion/sd3/sd3_medium_incl_clips.safetensors' - import time - import logging - logging.basicConfig(level=logging.INFO) - log = logging.getLogger('sd') - t0 = time.time() - pipeline = load_sd3(te3='fp16', fn='') - - # pipeline.to('cuda') - t1 = time.time() - log.info(f'Loaded: time={t1-t0:.3f}') - - # pipeline.scheduler = diffusers.schedulers.EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config) - log.info(f'Scheduler, {pipeline.scheduler}') - image = pipeline( - prompt='a photo of a cute robot holding a sign above his head that says sdnext, high detailed', - negative_prompt='', - num_inference_steps=50, - height=1024, - width=1024, - guidance_scale=7.0, - ).images[0] - t2 = time.time() - log.info(f'Generated: time={t2-t1:.3f}') - image.save("/tmp/sd3.png") + devices.torch_gc()