Views
No views yet

![FLUX.1 [schnell] Grid](./PEA-Diffusion.png)
MultilingualFLUX.1-adapter is a multilingual adapter tailored for the Flux.1 series models, theoretically, it inherits ByT5 and can support over 100 languages, but with additional optimizations in Chinese. Originating from an ECCV 2024 paper titled PEA-Diffusion. The open-source code is available at https://github.com/OPPO-Mente-Lab/PEA-Diffusion.MultilingualFLUX.1num_inference_steps and guidance_scale as needed.1from diffusers import FluxPipeline, AutoencoderKL
2from diffusers.image_processor import VaeImageProcessor
3from transformers import T5ForConditionalGeneration,AutoTokenizer
4import torch
5import torch.nn as nn
6
7
8class MLP(nn.Module):
9 def __init__(self, in_dim=4096, out_dim=4096, hidden_dim=4096, out_dim1=768, use_residual=True):
10 super().__init__()
11 self.layernorm = nn.LayerNorm(in_dim)
12 self.projector = nn.Sequential(
13 nn.Linear(in_dim, hidden_dim, bias=False),
14 nn.GELU(),
15 nn.Linear(hidden_dim, hidden_dim, bias=False),
16 nn.GELU(),
17 nn.Linear(hidden_dim, out_dim, bias=False),
18 )
19 self.fc = nn.Linear(out_dim, out_dim1)
20 def forward(self, x):
21 x = self.layernorm(x)
22 x = self.projector(x)
23 x2 = nn.GELU()(x)
24 x1 = self.fc(x2)
25 x1 = torch.mean(x1,1)
26 return x1,x2
27
28
29dtype = torch.bfloat16
30device = "cuda"
31ckpt_id = "black-forest-labs/FLUX.1-schnell"
32text_encoder_ckpt_id = 'google/byt5-xxl'
33proj_t5 = MLP(in_dim=4672, out_dim=4096, hidden_dim=4096, out_dim1=768).to(device=device,dtype=dtype)
34text_encoder_t5 = T5ForConditionalGeneration.from_pretrained(text_encoder_ckpt_id).get_encoder().to(device=device,dtype=dtype)
35tokenizer_t5 = AutoTokenizer.from_pretrained(text_encoder_ckpt_id)
36
37
38proj_t5_save_path = f"diffusion_pytorch_model.bin"
39state_dict = torch.load(proj_t5_save_path, map_location="cpu")
40state_dict_new = {}
41for k,v in state_dict.items():
42 k_new = k.replace("module.","")
43 state_dict_new[k_new] = v
44
45proj_t5.load_state_dict(state_dict_new)
46
47pipeline = FluxPipeline.from_pretrained(
48 ckpt_id, text_encoder=None, text_encoder_2=None,
49 tokenizer=None, tokenizer_2=None, vae=None,
50 torch_dtype=torch.bfloat16
51).to(device)
52
53vae = AutoencoderKL.from_pretrained(
54 ckpt_id,
55 subfolder="vae",
56 torch_dtype=torch.bfloat16
57).to(device)
58vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
59image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
60
61while True:
62 raw_text = input("\nPlease Input Query (stop to exit) >>> ")
63 if not raw_text:
64 print('Query should not be empty!')
65 continue
66 if raw_text == "stop":
67 break
68
69 with torch.no_grad():
70 text_inputs = tokenizer_t5(
71 raw_text,
72 padding="max_length",
73 max_length=256,
74 truncation=True,
75 add_special_tokens=True,
76 return_tensors="pt",
77 ).input_ids.to(device)
78 text_embeddings = text_encoder_t5(text_inputs)[0]
79 pooled_prompt_embeds,prompt_embeds = proj_t5(text_embeddings)
80 height, width = 1024, 1024
81 latents = pipeline(
82 prompt_embeds=prompt_embeds,
83 pooled_prompt_embeds=pooled_prompt_embeds,
84 num_inference_steps=4, guidance_scale=0,
85 height=height, width=width,
86 output_type="latent",
87 ).images
88
89 latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
90 latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
91 image = vae.decode(latents, return_dict=False)[0]
92 image = image_processor.postprocess(image, output_type="pil")
93 image[0].save("ChineseFLUX.jpg")
94MultilingualOpenFLUX.11from diffusers import FluxPipeline, AutoencoderKL
2from diffusers.image_processor import VaeImageProcessor
3from transformers import T5ForConditionalGeneration,AutoTokenizer
4import torch
5import torch.nn as nn
6
7
8class MLP(nn.Module):
9 def __init__(self, in_dim=4096, out_dim=4096, hidden_dim=4096, out_dim1=768, use_residual=True):
10 super().__init__()
11 self.layernorm = nn.LayerNorm(in_dim)
12 self.projector = nn.Sequential(
13 nn.Linear(in_dim, hidden_dim, bias=False),
14 nn.GELU(),
15 nn.Linear(hidden_dim, hidden_dim, bias=False),
16 nn.GELU(),
17 nn.Linear(hidden_dim, out_dim, bias=False),
18 )
19 self.fc = nn.Linear(out_dim, out_dim1)
20 def forward(self, x):
21 x = self.layernorm(x)
22 x = self.projector(x)
23 x2 = nn.GELU()(x)
24 x1 = self.fc(x2)
25 x1 = torch.mean(x1,1)
26 return x1,x2
27
28
29dtype = torch.bfloat16
30device = "cuda"
31ckpt_id = "ostris/OpenFLUX.1"
32text_encoder_ckpt_id = 'google/byt5-xxl'
33proj_t5 = MLP(in_dim=4672, out_dim=4096, hidden_dim=4096, out_dim1=768).to(device=device,dtype=dtype)
34text_encoder_t5 = T5ForConditionalGeneration.from_pretrained(text_encoder_ckpt_id).get_encoder().to(device=device,dtype=dtype)
35tokenizer_t5 = AutoTokenizer.from_pretrained(text_encoder_ckpt_id)
36
37
38proj_t5_save_path = f"diffusion_pytorch_model.bin"
39state_dict = torch.load(proj_t5_save_path, map_location="cpu")
40state_dict_new = {}
41for k,v in state_dict.items():
42 k_new = k.replace("module.","")
43 state_dict_new[k_new] = v
44
45proj_t5.load_state_dict(state_dict_new)
46
47pipeline = FluxPipeline.from_pretrained(
48 ckpt_id, text_encoder=None, text_encoder_2=None,
49 tokenizer=None, tokenizer_2=None, vae=None,
50 torch_dtype=torch.bfloat16
51).to(device)
52pipeline.load_lora_weights("ostris/OpenFLUX.1/openflux1-v0.1.0-fast-lora.safetensors")
53
54vae = AutoencoderKL.from_pretrained(
55 ckpt_id,
56 subfolder="vae",
57 torch_dtype=torch.bfloat16
58).to(device)
59vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
60image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
61
62while True:
63 raw_text = input("\nPlease Input Query (stop to exit) >>> ")
64 if not raw_text:
65 print('Query should not be empty!')
66 continue
67 if raw_text == "stop":
68 break
69
70 with torch.no_grad():
71 text_inputs = tokenizer_t5(
72 raw_text,
73 padding="max_length",
74 max_length=256,
75 truncation=True,
76 add_special_tokens=True,
77 return_tensors="pt",
78 ).input_ids.to(device)
79 text_embeddings = text_encoder_t5(text_inputs)[0]
80 pooled_prompt_embeds,prompt_embeds = proj_t5(text_embeddings)
81 height, width = 1024, 1024
82 latents = pipeline(
83 prompt_embeds=prompt_embeds,
84 pooled_prompt_embeds=pooled_prompt_embeds,
85 num_inference_steps=4, guidance_scale=0,
86 height=height, width=width,
87 output_type="latent",
88 ).images
89
90 latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
91 latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
92 image = vae.decode(latents, return_dict=False)[0]
93 image = image_processor.postprocess(image, output_type="pil")
94 image[0].save("ChineseOpenFLUX.jpg")
95