Skip to content

override Flux transformer #639

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Nerogar merged 16 commits into from
Apr 21, 2025
Merged

override Flux transformer #639

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
a40c25e
override Flux transformer
dxqb Jan 12, 2025
c5c08a2
ignore hugging face connection error for offline training
dxqb Jan 29, 2025
b60968d
pre-commit fix
dxqb Jan 29, 2025
395d078
Merge branch 'Nerogar:master' into master
dxqb Feb 2, 2025
a441f04
Merge branch 'Nerogar:master' into master
dxqb Feb 9, 2025
0fb8512
Merge branch 'Nerogar:master' into master
dxqb Feb 13, 2025
a1a263d
Merge branch 'Nerogar:master' into master
dxqb Feb 15, 2025
ae2e7ce
Merge branch 'master' of https://github.com/dxqbYD/OneTrainer
dxqb Feb 22, 2025
aec59b6
clean exit
dxqb Feb 22, 2025
15967e6
Merge branch 'Nerogar:master' into override_transformer
dxqb Feb 28, 2025
93fe9e0
Merge branch 'master' of https://github.com/dxqbYD/OneTrainer
dxqb Mar 8, 2025
6fefde0
Merge branch 'master' of https://github.com/dxqbYD/OneTrainer
dxqb Mar 8, 2025
af3cea8
Merge branch 'master' of https://github.com/dxqbYD/OneTrainer
dxqb Mar 16, 2025
0aa6c63
Merge branch 'master' of https://github.com/dxqbYD/OneTrainer
dxqb Mar 28, 2025
7929cb3
Merge branch 'master' into override_transformer
dxqb Mar 29, 2025
bf38e55
avoid loading in float32; bugfix
dxqb Mar 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions modules/cloud/BaseCloud.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ def upload_config(self,commands : TrainCommands=None):

if hasattr(self.config,"local_base_model_name"):
self.file_sync.sync_up(local=Path(self.config.local_base_model_name),remote=Path(self.config.base_model_name))
if hasattr(self.config.prior,"local_model_name"):
self.file_sync.sync_up(local=Path(self.config.prior.local_model_name),remote=Path(self.config.prior.model_name))
if hasattr(self.config,"local_lora_model_name"):
self.file_sync.sync_up(local=Path(self.config.local_lora_model_name),remote=Path(self.config.lora_model_name))

Expand Down
43 changes: 32 additions & 11 deletions modules/modelLoader/flux/FluxModelLoader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
from modules.util.ModelNames import ModelNames
from modules.util.ModelWeightDtypes import ModelWeightDtypes

import torch

from diffusers import (
AutoencoderKL,
FlowMatchEulerDiscreteScheduler,
Expand All @@ -28,13 +30,14 @@ def __load_internal(
model_type: ModelType,
weight_dtypes: ModelWeightDtypes,
base_model_name: str,
transformer_model_name: str,
vae_model_name: str,
include_text_encoder_1: bool,
include_text_encoder_2: bool,
):
if os.path.isfile(os.path.join(base_model_name, "meta.json")):
self.__load_diffusers(
model, model_type, weight_dtypes, base_model_name, vae_model_name,
model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name,
include_text_encoder_1, include_text_encoder_2,
)
else:
Expand All @@ -46,6 +49,7 @@ def __load_diffusers(
model_type: ModelType,
weight_dtypes: ModelWeightDtypes,
base_model_name: str,
transformer_model_name: str,
vae_model_name: str,
include_text_encoder_1: bool,
include_text_encoder_2: bool,
Expand Down Expand Up @@ -109,13 +113,23 @@ def __load_diffusers(
"vae",
)

transformer = self._load_diffusers_sub_module(
FluxTransformer2DModel,
weight_dtypes.prior,
weight_dtypes.train_dtype,
base_model_name,
"transformer",
)
if transformer_model_name:
transformer = FluxTransformer2DModel.from_single_file(
transformer_model_name,
#avoid loading the transformer in float32:
torch_dtype = torch.bfloat16 if weight_dtypes.prior.torch_dtype() is None else weight_dtypes.prior.torch_dtype()

@Nerogar Nerogar Apr 11, 2025 •
edited
Loading

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This could now quantize the weights twice if the source is not in bf16 format. For example, if weight_dtypes.prior is set to fp8, and the source is in fp32, then this will first convert to bf16 (losing precision). Then the actual quantization to fp8 is done.

@dxqb dxqb Apr 11, 2025 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is model-specific code and Flux was published in bf16, so I don't think we have to worry about precision loss from fp32.
the alternative is mentioned here, #639 (comment), to let the user choose a load dtype.

The current behaviour is to load a fp8 source to fp32 and load a bf16 source to bf16. Loading fp8 sources needs 48 GB vram just for that.
new behaviour is to load both to bf16.

@Nerogar Nerogar Apr 12, 2025 •
edited
Loading

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But that argument is only true for the base flux model released by BFL. Fine tunes can be distributed in any other precision. There was some discussion about a possible solution in the diffusers library, but it didn't really reach any conclusion huggingface/diffusers#10679

A load d type could work, but I can almost guarantee you that there will be a lot of questions about that setting

@dxqb dxqb Apr 18, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd leave it as proposed in the PR.
It's not ideal, but an improvement for all current use cases, which is fp8 and bf16 transformer files. Precision loss can happen only if you had a fp32 transformer file and wanted to load it in fp8, and what's the point of that? If you specifically created an fp32 transformer for some reason on your own (because none exist AFAIK), you would load it in fp32 and not in fp8

)
transformer = self._convert_diffusers_sub_module_to_dtype(
transformer, weight_dtypes.prior, weight_dtypes.train_dtype
)
else:
transformer = self._load_diffusers_sub_module(
FluxTransformer2DModel,
weight_dtypes.prior,
weight_dtypes.train_dtype,
base_model_name,
"transformer",
)

model.model_type = model_type
model.tokenizer_1 = tokenizer_1
Expand Down Expand Up @@ -146,13 +160,20 @@ def __load_safetensors(
model_type: ModelType,
weight_dtypes: ModelWeightDtypes,
base_model_name: str,
transformer_model_name: str,
vae_model_name: str,
include_text_encoder_1: bool,
include_text_encoder_2: bool,
):
transformer = FluxTransformer2DModel.from_single_file(
#always load transformer separately even though FluxPipeLine.from_single_file() could load it, to avoid loading in float32:
transformer_model_name if transformer_model_name else base_model_name,
torch_dtype = torch.bfloat16 if weight_dtypes.prior.torch_dtype() is None else weight_dtypes.prior.torch_dtype()
)
pipeline = FluxPipeline.from_single_file(
pretrained_model_link_or_path=base_model_name,
safety_checker=None,
transformer=transformer,
Comment thread
dxqb marked this conversation as resolved.
Resolved
)

if include_text_encoder_2:
Expand Down Expand Up @@ -218,7 +239,7 @@ def load(

try:
self.__load_internal(
model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model,
model, model_type, weight_dtypes, model_names.base_model, model_names.prior_model, model_names.vae_model,
model_names.include_text_encoder, model_names.include_text_encoder_2,
)
return
Expand All @@ -227,7 +248,7 @@ def load(

try:
self.__load_diffusers(
model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model,
model, model_type, weight_dtypes, model_names.base_model, model_names.prior_model, model_names.vae_model,
model_names.include_text_encoder, model_names.include_text_encoder_2,
)
return
Expand All @@ -236,7 +257,7 @@ def load(

try:
self.__load_safetensors(
model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model,
model, model_type, weight_dtypes, model_names.base_model, model_names.prior_model, model_names.vae_model,
model_names.include_text_encoder, model_names.include_text_encoder_2,
)
return
Expand Down
2 changes: 2 additions & 0 deletions modules/trainer/CloudTrainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,8 @@ def adjust(config,attribute : str):
adjust(remote,"cache_dir")
if Path(remote.base_model_name).exists() or remote.base_model_name.startswith("cloud:"):
adjust(remote,"base_model_name")
if Path(remote.prior.model_name).exists() or remote.prior.model_name.startswith("cloud:"):
adjust(remote.prior,"model_name")

adjust(remote,"output_model_destination")
adjust(remote,"lora_model_name")
Expand Down
3 changes: 2 additions & 1 deletion modules/ui/ModelTab.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ def __setup_flux_ui(self):
row = self.__create_base_components(
row,
has_prior=True,
allow_override_prior=True,
has_text_encoder_1=True,
has_text_encoder_2=True,
has_vae=True,
Expand Down Expand Up @@ -265,7 +266,7 @@ def __create_base_components(
if allow_override_prior:
# prior model
components.label(self.scroll_frame, row, 0, "Prior Model",
tooltip="Filename, directory or Hugging Face repository of the prior model")
tooltip="Filename, directory or Hugging Face repository of the prior model. For Flux, it must be a safetensors file that only contains the transformer.")
Comment thread
dxqb marked this conversation as resolved.
Resolved
components.file_entry(
self.scroll_frame, row, 1, self.ui_state, "prior.model_name",
path_modifier=lambda x: Path(x).parent.absolute() if x.endswith(".json") else x
Expand Down