iframe-proxy | Sunbelt Computer Software

Repository navigation

History

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

import argparse

import json

import os

import safetensors.torch

import torch

from accelerate import init_empty_weights

from huggingface_hub import snapshot_download

from transformers import AutoImageProcessor, Qwen3VLForConditionalGeneration, Qwen3VLProcessor

from diffusers import (

AutoencoderKLFlux2,

BriaFibo2Pipeline,

BriaFibo2Transformer2DModel,

FlowMatchEulerDiscreteScheduler,

)

from diffusers.configuration_utils import FrozenDict

"""

Converts an original Fibo 2.0 checkpoint, with `transformer/`, `vae/`, `text_encoder/` and `scheduler/` folders, to

diffusers. The converted checkpoints are released as briaai/Fibo-2.0-dev and briaai/Fibo-2.0-edit-turbo.

# Transformer only

python scripts/convert_bria_fibo2_to_diffusers.py \

--checkpoint_path /path/to/original/checkpoint \

--output_path Fibo-2.0-dev

# Full pipeline

python scripts/convert_bria_fibo2_to_diffusers.py \

--checkpoint_path /path/to/original/checkpoint \

--output_path Fibo-2.0-dev \

--full_pipe

"""

parser = argparse.ArgumentParser()

parser.add_argument("--original_repo_id", default=None, type=str)

parser.add_argument("--checkpoint_path", default=None, type=str)

parser.add_argument("--output_path", type=str, required=True)

parser.add_argument("--full_pipe", action="store_true")

args = parser.parse_args()

TRANSFORMER_KEYS_RENAME_DICT = {

# one embedder and one final layer instead of Z-Image's one per patch size

"all_x_embedder.2-1.": "x_embedder.",

"all_final_layer.2-1.": "final_layer.",

# injection blocks and Perceiver layers: cross-attention to the text encoder

"cross_q.": "cross_attention.to_q.",

"cross_k.": "cross_attention.to_k.",

"cross_v.": "cross_attention.to_v.",

"cross_out.": "cross_attention.to_out.0.",

"cross_norm_q.": "cross_attention.norm_q.",

"cross_norm_k.": "cross_attention.norm_k.",

# Perceiver layers: self-attention among the gist tokens

"self_q.": "self_attention.to_q.",

"self_k.": "self_attention.to_k.",

"self_v.": "self_attention.to_v.",

"self_out.": "self_attention.to_out.0.",

"self_norm_q.": "self_attention.norm_q.",

"self_norm_k.": "self_attention.norm_k.",

}

# BriaFibo2Pipeline.text_encoder_out_layers: the Qwen3-VL layers behind the Perceiver, then each injection block

PIPELINE_TEXT_ENCODER_LAYERS = [[9, 20, 31], [5, 16, 27], [7, 18, 29], [9, 20, 31], [11, 22, 33], [13, 24, 35]]

def get_checkpoint_path(args):

if args.original_repo_id is not None:

allow_patterns = None if args.full_pipe else ["transformer/*"]

return snapshot_download(args.original_repo_id, allow_patterns=allow_patterns)

elif args.checkpoint_path is not None:

return args.checkpoint_path

else:

raise ValueError("please provide either `original_repo_id` or a local `checkpoint_path`")

def load_original_checkpoint(checkpoint_path):

transformer_path = os.path.join(checkpoint_path, "transformer")

with open(os.path.join(transformer_path, "config.json")) as f:

original_config = json.load(f)

with open(os.path.join(transformer_path, "diffusion_pytorch_model.safetensors.index.json")) as f:

shard_names = sorted(set(json.load(f)["weight_map"].values()))

original_state_dict = {}

for shard_name in shard_names:

original_state_dict.update(safetensors.torch.load_file(os.path.join(transformer_path, shard_name)))

return original_config, original_state_dict

def get_transformer_config(original_config):

if original_config["perceiver_num_heads"] not in (None, original_config["n_heads"]):

raise ValueError("The Perceiver must use as many attention heads as the main blocks.")

# BriaFibo2Pipeline feeds the transformer fixed Qwen3-VL layers: the Perceiver's, then three per injection block

text_encoder_layers = [original_config["perceiver_llm_layers"], *original_config["injection_llm_layers"]]

if text_encoder_layers != PIPELINE_TEXT_ENCODER_LAYERS:

raise ValueError(

f"The checkpoint reads Qwen3-VL layers {text_encoder_layers}, but BriaFibo2Pipeline feeds it"

f" {PIPELINE_TEXT_ENCODER_LAYERS}."

)

return {

"in_channels": original_config["in_channels"],

"patch_size": original_config["patch_size"],

"dim": original_config["dim"],

"n_layers": original_config["n_layers"],

"n_refiner_layers": original_config["n_refiner_layers"],

"n_heads": original_config["n_heads"],

"cap_feat_dim": original_config["cond_llm_dim"],

"injection_layer_ids": original_config["injection_layer_ids"],

"perceiver_num_layers": original_config["perceiver_num_layers"],

"min_num_gist_tokens": original_config["min_num_gist_tokens"],

"max_num_gist_tokens": original_config["max_num_gist_tokens"],

"gist_step": original_config["gist_step"],

"gist_min_text_len": original_config["gist_min_text_len"],

"gist_max_text_len": original_config["gist_max_text_len"],

}

def convert_bria_fibo2_transformer_to_diffusers(original_config, original_state_dict):

with init_empty_weights():

transformer = BriaFibo2Transformer2DModel.from_config(get_transformer_config(original_config))

converted_state_dict = {}

for key, value in original_state_dict.items():

new_key = key

for old, new in TRANSFORMER_KEYS_RENAME_DICT.items():

new_key = new_key.replace(old, new)

converted_state_dict[new_key] = value

transformer.load_state_dict(converted_state_dict, strict=True, assign=True)

return transformer

def main(args):

checkpoint_path = get_checkpoint_path(args)

original_config, original_state_dict = load_original_checkpoint(checkpoint_path)

transformer = convert_bria_fibo2_transformer_to_diffusers(original_config, original_state_dict)

if not args.full_pipe:

transformer.save_pretrained(os.path.join(args.output_path, "transformer"))

return

vae = AutoencoderKLFlux2.from_pretrained(checkpoint_path, subfolder="vae")

# the original config carries `force_upcast`, which the Flux 2 VAE doesn't take; left in, every load warns about it

vae_config = dict(vae.config)

vae_config.pop("force_upcast", None)

vae._internal_dict = FrozenDict(vae_config)

text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(

checkpoint_path, subfolder="text_encoder", dtype=torch.bfloat16

)

# Fibo 2.0's limits on an image's size in pixels, not Qwen3-VL's defaults; they decide how many tokens an image to

# edit becomes

processor = Qwen3VLProcessor.from_pretrained(checkpoint_path, subfolder="text_encoder")

processor.image_processor = AutoImageProcessor.from_pretrained(

checkpoint_path, subfolder="text_encoder", size={"shortest_edge": 4 * 28 * 28, "longest_edge": 2048 * 28 * 28}

)

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(checkpoint_path, subfolder="scheduler")

pipe = BriaFibo2Pipeline(

transformer=transformer, scheduler=scheduler, vae=vae, text_encoder=text_encoder, processor=processor

)

pipe.save_pretrained(args.output_path)

if __name__ == "__main__":

main(args)

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Repository navigation

convert_bria_fibo2_to_diffusers.py

convert_bria_fibo2_to_diffusers.py

Sunbelt Computer Software

PL/B Language Development and Support

FilesExpand file tree

convert_bria_fibo2_to_diffusers.py

Latest commit

History

convert_bria_fibo2_to_diffusers.py

File metadata and controls