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

sd3 simplify safetensors handler

Browse Source
This commit is contained in:
Vladimir Mandic
2024-06-13 19:10:29 -04:00
parent 3195b0fd7b
commit c31b19cb57
1 changed files with 44 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_sd3.py
+44 -97
View File
@@ -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()