CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-06-04 21:05:48 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fp32 support

Browse Source
This commit is contained in:
maybleMyers 2025-11-27 12:55:10 -08:00
parent b83ddf9e82
commit d52af5ee37
3 changed files with 156 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
backend/diffusion_engine/zimage.py
View File

@ -128,6 +128,13 @@ class ZImage(ForgeDiffusionEngine):
print(f"Transformed timestep: {original_timestep} -> {timestep}")
print(f"(Official would use: (1000 - t)/1000 where t ≈ sigma*1000)")
# Ensure pad tokens are on the same device as input (fixes CPU/GPU mismatch when model is partially offloaded)
target_device = x[0].device
if hasattr(self.transformer, 'x_pad_token') and self.transformer.x_pad_token.device != target_device:
self.transformer.x_pad_token.data = self.transformer.x_pad_token.data.to(target_device)
if hasattr(self.transformer, 'cap_pad_token') and self.transformer.cap_pad_token.device != target_device:
self.transformer.cap_pad_token.data = self.transformer.cap_pad_token.data.to(target_device)
# Call transformer with correct format
# patch_size=2 and f_patch_size=1 are the only supported values per config
# Transformer returns (list_of_tensors, {}), extract the list

76
backend/loader.py
View File

@ -54,7 +54,25 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
config = IntegratedAutoencoderKL.load_config(config_path)
with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype()):
# Check for Z-Image specific VAE precision
vae_dtype = memory_management.vae_dtype()
if getattr(guess, 'is_zimage', False):
try:
from modules import shared
z_vae_dtype = getattr(shared.opts, 'z_vae_dtype', 'Automatic')
if z_vae_dtype != 'Automatic':
dtype_map = {
'bfloat16': torch.bfloat16,
'float16': torch.float16,
'float32': torch.float32,
}
if z_vae_dtype in dtype_map:
vae_dtype = dtype_map[z_vae_dtype]
print(f'Z-Image VAE: Using user-specified dtype: {vae_dtype}')
except Exception as e:
print(f'Warning: Could not read Z-Image VAE precision setting: {e}')
with using_forge_operations(device=memory_management.cpu, dtype=vae_dtype):
model = IntegratedAutoencoderKL.from_config(config)
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in state_dict.keys(): #diffusers format
@ -113,8 +131,27 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
return model
if cls_name == 'Qwen3Model':
# Load Qwen model directly from transformers
text_encoder_dtype = memory_management.text_encoder_dtype()
# Check for Z-Image specific text encoder precision
if getattr(guess, 'is_zimage', False):
try:
from modules import shared
z_text_encoder_dtype = getattr(shared.opts, 'z_text_encoder_dtype', 'Automatic')
if z_text_encoder_dtype != 'Automatic':
dtype_map = {
'bfloat16': torch.bfloat16,
'float16': torch.float16,
'float32': torch.float32,
}
if z_text_encoder_dtype in dtype_map:
text_encoder_dtype = dtype_map[z_text_encoder_dtype]
print(f'Z-Image Text Encoder: Using user-specified dtype: {text_encoder_dtype}')
except Exception as e:
print(f'Warning: Could not read Z-Image text encoder precision setting: {e}')
cls = getattr(importlib.import_module('transformers'), cls_name)
model = cls.from_pretrained(config_path, torch_dtype=memory_management.text_encoder_dtype())
model = cls.from_pretrained(config_path, torch_dtype=text_encoder_dtype)
return model
if cls_name in ['UNet2DConditionModel', 'FluxTransformer2DModel', 'SD3Transformer2DModel', 'ChromaTransformer2DModel', 'ChromaDCTTransformer2DModel', 'ZImageTransformer2DModel']:
assert isinstance(state_dict, dict) and len(state_dict) > 16, 'You do not have model state dict!'
@ -157,8 +194,42 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
if state_dict_dtype in ['gguf']:
beautiful_print_gguf_state_dict_statics(state_dict)
# Z-Image specific precision settings (highest priority)
if cls_name == 'ZImageTransformer2DModel':
try:
from modules import shared
z_transformer_dtype = getattr(shared.opts, 'z_transformer_dtype', 'Automatic')
if z_transformer_dtype != 'Automatic':
dtype_map = {
'bfloat16': torch.bfloat16,
'float16': torch.float16,
'float32': torch.float32,
}
if z_transformer_dtype in dtype_map:
storage_dtype = dtype_map[z_transformer_dtype]
print(f'Z-Image Transformer: Using user-specified dtype: {storage_dtype}')
except Exception as e:
print(f'Warning: Could not read Z-Image precision setting: {e}')
load_device = memory_management.get_torch_device()
computation_dtype = memory_management.get_computation_dtype(load_device, parameters=state_dict_parameters, supported_dtypes=guess.supported_inference_dtypes)
# For Z-Image, if user specified a precision, also use it for computation
if cls_name == 'ZImageTransformer2DModel':
try:
from modules import shared
z_transformer_dtype = getattr(shared.opts, 'z_transformer_dtype', 'Automatic')
if z_transformer_dtype != 'Automatic':
dtype_map = {
'bfloat16': torch.bfloat16,
'float16': torch.float16,
'float32': torch.float32,
}
if z_transformer_dtype in dtype_map:
computation_dtype = dtype_map[z_transformer_dtype]
except Exception:
pass
offload_device = memory_management.unet_offload_device()
if storage_dtype in ['nf4', 'fp4', 'gguf']:
@ -569,6 +640,7 @@ def forge_loader(sd, additional_state_dicts=None, preset=None):
class ZImageConfig:
def __init__(self):
self.huggingface_repo = "Z-Image"
self.is_zimage = True # Flag to identify Z-Image models
self.unet_config = {
'in_channels': 16,
'dim': 3840,

90
modules_forge/main_entry.py
View File

@ -23,6 +23,11 @@ ui_forge_async_loading: gr.Radio = None
ui_forge_pin_shared_memory: gr.Radio = None
ui_forge_inference_memory: gr.Slider = None
# Z-Image specific precision options
ui_z_transformer_dtype: gr.Dropdown = None
ui_z_vae_dtype: gr.Dropdown = None
ui_z_text_encoder_dtype: gr.Dropdown = None
forge_unet_storage_dtype_options = {
@ -38,6 +43,9 @@ forge_unet_storage_dtype_options = {
'float8-e5m2 (fp16 LoRA)': (torch.float8_e5m2, True),
}
# Z-Image precision options
z_precision_options = ['Automatic', 'bfloat16', 'float16', 'float32']
module_list = {}
@ -55,7 +63,7 @@ def bind_to_opts(comp, k, save=False, callback=None):
def make_checkpoint_manager_ui():
global ui_checkpoint, ui_vae, ui_clip_skip, ui_forge_unet_storage_dtype_options, ui_forge_async_loading, ui_forge_pin_shared_memory, ui_forge_inference_memory, ui_forge_preset
global ui_checkpoint, ui_vae, ui_clip_skip, ui_forge_unet_storage_dtype_options, ui_forge_async_loading, ui_forge_pin_shared_memory, ui_forge_inference_memory, ui_forge_preset, ui_z_transformer_dtype, ui_z_vae_dtype, ui_z_text_encoder_dtype
if shared.opts.sd_model_checkpoint in [None, 'None', 'none', '']:
if len(sd_models.checkpoints_list) == 0:
@ -122,6 +130,16 @@ def make_checkpoint_manager_ui():
ui_clip_skip = gr.Slider(label="Clip skip", value=lambda: shared.opts.CLIP_stop_at_last_layers, **{"minimum": 1, "maximum": 12, "step": 1})
bind_to_opts(ui_clip_skip, 'CLIP_stop_at_last_layers', save=True)
# Z-Image precision options
ui_z_transformer_dtype = gr.Dropdown(label="Z-Image Transformer Precision", value=lambda: shared.opts.z_transformer_dtype, choices=z_precision_options)
bind_to_opts(ui_z_transformer_dtype, 'z_transformer_dtype', save=True, callback=refresh_model_loading_parameters)
ui_z_vae_dtype = gr.Dropdown(label="Z-Image VAE Precision", value=lambda: shared.opts.z_vae_dtype, choices=z_precision_options)
bind_to_opts(ui_z_vae_dtype, 'z_vae_dtype', save=True, callback=refresh_model_loading_parameters)
ui_z_text_encoder_dtype = gr.Dropdown(label="Z-Image Text Encoder Precision", value=lambda: shared.opts.z_text_encoder_dtype, choices=z_precision_options)
bind_to_opts(ui_z_text_encoder_dtype, 'z_text_encoder_dtype', save=True, callback=refresh_model_loading_parameters)
ui_checkpoint.change(checkpoint_change, inputs=[ui_checkpoint], show_progress=False)
ui_vae.change(modules_change, inputs=[ui_vae], queue=False, show_progress=False)
@ -294,6 +312,7 @@ def forge_main_entry():
ui_txt2img_distilled_cfg = get_a1111_ui_component('txt2img', 'Distilled CFG Scale')
ui_txt2img_sampler = get_a1111_ui_component('txt2img', 'sampler_name')
ui_txt2img_scheduler = get_a1111_ui_component('txt2img', 'scheduler')
ui_txt2img_steps = get_a1111_ui_component('txt2img', 'steps')
ui_img2img_width = get_a1111_ui_component('img2img', 'Size-1')
ui_img2img_height = get_a1111_ui_component('img2img', 'Size-2')
@ -301,6 +320,7 @@ def forge_main_entry():
ui_img2img_distilled_cfg = get_a1111_ui_component('img2img', 'Distilled CFG Scale')
ui_img2img_sampler = get_a1111_ui_component('img2img', 'sampler_name')
ui_img2img_scheduler = get_a1111_ui_component('img2img', 'scheduler')
ui_img2img_steps = get_a1111_ui_component('img2img', 'steps')
ui_txt2img_hr_cfg = get_a1111_ui_component('txt2img', 'Hires CFG Scale')
ui_txt2img_hr_distilled_cfg = get_a1111_ui_component('txt2img', 'Hires Distilled CFG Scale')
@ -326,6 +346,11 @@ def forge_main_entry():
ui_img2img_scheduler,
ui_txt2img_hr_cfg,
ui_txt2img_hr_distilled_cfg,
ui_txt2img_steps,
ui_img2img_steps,
ui_z_transformer_dtype,
ui_z_vae_dtype,
ui_z_text_encoder_dtype,
]
ui_forge_preset.change(on_preset_change, inputs=[ui_forge_preset], outputs=output_targets, queue=False, show_progress=False)
@ -365,6 +390,11 @@ def on_preset_change(preset=None):
gr.update(value=getattr(shared.opts, "sd_i2i_scheduler", 'Automatic')), # ui_img2img_scheduler
gr.update(visible=True, value=getattr(shared.opts, "sd_t2i_hr_cfg", 7.0)), # ui_txt2img_hr_cfg
gr.update(visible=False, value=3.5), # ui_txt2img_hr_distilled_cfg
gr.update(value=20), # ui_txt2img_steps
gr.update(value=20), # ui_img2img_steps
gr.update(visible=False), # ui_z_transformer_dtype
gr.update(visible=False), # ui_z_vae_dtype
gr.update(visible=False), # ui_z_text_encoder_dtype
]
if shared.opts.forge_preset == 'xl':
@ -394,6 +424,11 @@ def on_preset_change(preset=None):
gr.update(value=getattr(shared.opts, "xl_i2i_scheduler", 'Automatic')), # ui_img2img_scheduler
gr.update(visible=True, value=getattr(shared.opts, "xl_t2i_hr_cfg", 5.0)), # ui_txt2img_hr_cfg
gr.update(visible=False, value=3.5), # ui_txt2img_hr_distilled_cfg
gr.update(value=20), # ui_txt2img_steps
gr.update(value=20), # ui_img2img_steps
gr.update(visible=False), # ui_z_transformer_dtype
gr.update(visible=False), # ui_z_vae_dtype
gr.update(visible=False), # ui_z_text_encoder_dtype
]
if shared.opts.forge_preset == 'flux':
@ -423,6 +458,11 @@ def on_preset_change(preset=None):
gr.update(value=getattr(shared.opts, "flux_i2i_scheduler", 'Simple')), # ui_img2img_scheduler
gr.update(visible=True, value=getattr(shared.opts, "flux_t2i_hr_cfg", 1.0)), # ui_txt2img_hr_cfg
gr.update(visible=True, value=getattr(shared.opts, "flux_t2i_hr_d_cfg", 3.5)), # ui_txt2img_hr_distilled_cfg
gr.update(value=20), # ui_txt2img_steps
gr.update(value=20), # ui_img2img_steps
gr.update(visible=False), # ui_z_transformer_dtype
gr.update(visible=False), # ui_z_vae_dtype
gr.update(visible=False), # ui_z_text_encoder_dtype
]
if shared.opts.forge_preset == 'chroma':
@ -449,6 +489,11 @@ def on_preset_change(preset=None):
gr.update(value=getattr(shared.opts, "chroma_i2i_scheduler", 'Simple')), # ui_img2img_scheduler
gr.update(visible=True, value=getattr(shared.opts, "chroma_t2i_hr_cfg", 1.0)), # ui_txt2img_hr_cfg
gr.update(visible=True, value=getattr(shared.opts, "chroma_t2i_hr_d_cfg", 3.5)), # ui_txt2img_hr_distilled_cfg
gr.update(value=20), # ui_txt2img_steps
gr.update(value=20), # ui_img2img_steps
gr.update(visible=False), # ui_z_transformer_dtype
gr.update(visible=False), # ui_z_vae_dtype
gr.update(visible=False), # ui_z_text_encoder_dtype
]
if shared.opts.forge_preset == 'z':
@ -464,20 +509,25 @@ def on_preset_change(preset=None):
gr.update(visible=True, value='Queue'), # ui_forge_async_loading
gr.update(visible=True, value='CPU'), # ui_forge_pin_shared_memory
gr.update(visible=True, value=model_mem), # ui_forge_inference_memory
gr.update(maximum=2048, value=getattr(shared.opts, "z_t2i_width", 1024)), # ui_txt2img_width
gr.update(maximum=2048, value=getattr(shared.opts, "z_i2i_width", 1024)), # ui_img2img_width
gr.update(maximum=2048, value=getattr(shared.opts, "z_t2i_height", 1024)), # ui_txt2img_height
gr.update(maximum=2048, value=getattr(shared.opts, "z_i2i_height", 1024)), # ui_img2img_height
gr.update(value=getattr(shared.opts, "z_t2i_cfg", 0)), # ui_txt2img_cfg
gr.update(value=getattr(shared.opts, "z_i2i_cfg", 0)), # ui_img2img_cfg
gr.update(maximum=3072, value=getattr(shared.opts, "z_t2i_width", 1024)), # ui_txt2img_width
gr.update(maximum=3072, value=getattr(shared.opts, "z_i2i_width", 1024)), # ui_img2img_width
gr.update(maximum=3072, value=getattr(shared.opts, "z_t2i_height", 1024)), # ui_txt2img_height
gr.update(maximum=3072, value=getattr(shared.opts, "z_i2i_height", 1024)), # ui_img2img_height
gr.update(value=getattr(shared.opts, "z_t2i_cfg", 1)), # ui_txt2img_cfg
gr.update(value=getattr(shared.opts, "z_i2i_cfg", 1)), # ui_img2img_cfg
gr.update(visible=False, value=0), # ui_txt2img_distilled_cfg
gr.update(visible=False, value=0), # ui_img2img_distilled_cfg
gr.update(value=getattr(shared.opts, "z_t2i_sampler", 'Euler')), # ui_txt2img_sampler
gr.update(value=getattr(shared.opts, "z_i2i_sampler", 'Euler')), # ui_img2img_sampler
gr.update(value=getattr(shared.opts, "z_t2i_scheduler", 'Simple')), # ui_txt2img_scheduler
gr.update(value=getattr(shared.opts, "z_i2i_scheduler", 'Simple')), # ui_img2img_scheduler
gr.update(visible=True, value=getattr(shared.opts, "z_t2i_hr_cfg", 0)), # ui_txt2img_hr_cfg
gr.update(visible=True, value=getattr(shared.opts, "z_t2i_hr_cfg", 1)), # ui_txt2img_hr_cfg
gr.update(visible=False, value=0), # ui_txt2img_hr_distilled_cfg
gr.update(value=getattr(shared.opts, "z_t2i_steps", 9)), # ui_txt2img_steps
gr.update(value=getattr(shared.opts, "z_i2i_steps", 9)), # ui_img2img_steps
gr.update(visible=True, value=getattr(shared.opts, "z_transformer_dtype", 'Automatic')), # ui_z_transformer_dtype
gr.update(visible=True, value=getattr(shared.opts, "z_vae_dtype", 'Automatic')), # ui_z_vae_dtype
gr.update(visible=True, value=getattr(shared.opts, "z_text_encoder_dtype", 'Automatic')), # ui_z_text_encoder_dtype
]
loadsave = ui_loadsave.UiLoadsave(cmd_opts.ui_config_file)
@ -504,6 +554,11 @@ def on_preset_change(preset=None):
gr.update(value=ui_settings_from_file['customscript/sampler.py/img2img/Schedule type/value']), # ui_img2img_scheduler
gr.update(visible=True, value=ui_settings_from_file['txt2img/Hires CFG Scale/value']), # ui_txt2img_hr_cfg
gr.update(visible=True, value=ui_settings_from_file['txt2img/Hires Distilled CFG Scale/value']), # ui_txt2img_hr_distilled_cfg
gr.update(value=ui_settings_from_file.get('customscript/sampler.py/txt2img/Sampling steps/value', 20)), # ui_txt2img_steps
gr.update(value=ui_settings_from_file.get('customscript/sampler.py/img2img/Sampling steps/value', 20)), # ui_img2img_steps
gr.update(visible=True, value=getattr(shared.opts, "z_transformer_dtype", 'Automatic')), # ui_z_transformer_dtype
gr.update(visible=True, value=getattr(shared.opts, "z_vae_dtype", 'Automatic')), # ui_z_vae_dtype
gr.update(visible=True, value=getattr(shared.opts, "z_text_encoder_dtype", 'Automatic')), # ui_z_text_encoder_dtype
]
shared.options_templates.update(shared.options_section(('ui_sd', "UI defaults 'sd'", "ui"), {
@ -556,16 +611,21 @@ shared.options_templates.update(shared.options_section(('ui_chroma', "UI default
"chroma_GPU_MB": shared.OptionInfo(total_vram - 1024, "GPU Weights (MB)",gr.Slider, {"minimum": 0, "maximum": total_vram, "step": 1}),
}))
shared.options_templates.update(shared.options_section(('ui_z', "UI defaults 'z' (Z-Image)", "ui"), {
"z_t2i_width": shared.OptionInfo(1024, "txt2img width", gr.Slider, {"minimum": 64, "maximum": 2048, "step": 8}),
"z_t2i_height": shared.OptionInfo(1024, "txt2img height", gr.Slider, {"minimum": 64, "maximum": 2048, "step": 8}),
"z_t2i_width": shared.OptionInfo(1024, "txt2img width", gr.Slider, {"minimum": 64, "maximum": 3072, "step": 8}),
"z_t2i_height": shared.OptionInfo(1024, "txt2img height", gr.Slider, {"minimum": 64, "maximum": 3072, "step": 8}),
"z_t2i_steps": shared.OptionInfo(9, "txt2img steps", gr.Slider, {"minimum": 1, "maximum": 150, "step": 1}),
"z_t2i_sampler": shared.OptionInfo('Euler', "txt2img sampler"),
"z_t2i_scheduler": shared.OptionInfo('Simple', "txt2img scheduler"),
"z_t2i_cfg": shared.OptionInfo(0, "txt2img CFG", gr.Slider, {"minimum": 0, "maximum": 30, "step": 0.1}),
"z_t2i_hr_cfg": shared.OptionInfo(0, "txt2img HiRes CFG", gr.Slider, {"minimum": 0, "maximum": 30, "step": 0.1}),
"z_i2i_width": shared.OptionInfo(1024, "img2img width", gr.Slider, {"minimum": 64, "maximum": 2048, "step": 8}),
"z_i2i_height": shared.OptionInfo(1024, "img2img height", gr.Slider, {"minimum": 64, "maximum": 2048, "step": 8}),
"z_t2i_cfg": shared.OptionInfo(1, "txt2img CFG", gr.Slider, {"minimum": 0, "maximum": 30, "step": 0.1}),
"z_t2i_hr_cfg": shared.OptionInfo(1, "txt2img HiRes CFG", gr.Slider, {"minimum": 0, "maximum": 30, "step": 0.1}),
"z_i2i_width": shared.OptionInfo(1024, "img2img width", gr.Slider, {"minimum": 64, "maximum": 3072, "step": 8}),
"z_i2i_height": shared.OptionInfo(1024, "img2img height", gr.Slider, {"minimum": 64, "maximum": 3072, "step": 8}),
"z_i2i_steps": shared.OptionInfo(9, "img2img steps", gr.Slider, {"minimum": 1, "maximum": 150, "step": 1}),
"z_i2i_sampler": shared.OptionInfo('Euler', "img2img sampler"),
"z_i2i_scheduler": shared.OptionInfo('Simple', "img2img scheduler"),
"z_i2i_cfg": shared.OptionInfo(0, "img2img CFG", gr.Slider, {"minimum": 0, "maximum": 30, "step": 0.1}),
"z_i2i_cfg": shared.OptionInfo(1, "img2img CFG", gr.Slider, {"minimum": 0, "maximum": 30, "step": 0.1}),
"z_GPU_MB": shared.OptionInfo(total_vram - 1024, "GPU Weights (MB)", gr.Slider, {"minimum": 0, "maximum": total_vram, "step": 1}),
"z_transformer_dtype": shared.OptionInfo('Automatic', "Transformer Precision", gr.Dropdown, {"choices": z_precision_options}),
"z_vae_dtype": shared.OptionInfo('Automatic', "VAE Precision", gr.Dropdown, {"choices": z_precision_options}),
"z_text_encoder_dtype": shared.OptionInfo('Automatic', "Text Encoder Precision", gr.Dropdown, {"choices": z_precision_options}),
}))