OPPO开源Diffusion多语言适配器—— MultilingualSD3-adapter 和 ChineseFL ...

MultilingualSD3-adapter 是为 SD3 量身定制的多语言适配器。 它源自 ECCV 2024 的一篇题为 PEA-Diffusion 的论文。ChineseFLUX.1-adapter是为Flux.1系列机型量身定制的多语言适配器，理论上继续了ByT5，可支持100多种语言，但在中文方面做了额外优化。 它源于一篇题为 PEA-Diffusion 的 ECCV 2024 论文。

PEA-Diffusion

作者首先提到了一种名为 “知识蒸馏”（Knowledge Distillation，KD）的方法。KD 是一个过程，在这个过程中，一个较小的机器学习模型（称为学生）会学习模仿一个较大、较复杂的模型（称为西席）。在这种情况下，目的是让学生模型尽可能接近西席的输出或预测。这被称为 “逼迫学生分布与西席分布相匹配”。
然而，作者还盼望确保西席模型和学生模型不仅能产生相似的输出效果，而且能以相似的方式处理信息。这就是特征对齐的作用所在。他们盼望西席模型和学生模型的中间计算或特征图也能相似。因此，他们引入了一种技术，在这些模型的中间层增强这种特征对齐。如许做的目的是减少所谓的分布偏移，使学生模型更像西席模型。
现在，当模型之间存在维度差异时，OPPOer/ChineseFLUX.1-adapter 就会发挥作用。在这种情况下，他们使用的是名为 CLIP（对比语言图像预练习）的模型，一个用于英语数据，一个用于非英语数据。适配器是添加到模型中的一个小型附加组件，用于对齐或转换特征（中间计算）以匹配维度。该适配器颠末练习，可以学习特定语言的信息，而且设计高效，只有 600 万个参数。
因此，总的来说，OPPOer/ChineseFLUX.1-adapter 是一种用于调整不同模型（在本例中为英语和非英语 CLIP 模型）中特征维度的技术，以促进知识提炼并进步学生模型的性能。它是这个复杂系统中一个小而紧张的组成部分！

论文：PEA-Diffusion: Parameter-Efficient Adapter with Knowledge Distillation in non-English Text-to-Image Generation
Code： https://github.com/OPPO-Mente-Lab/PEA-Diffusion
MultilingualSD3-adapter

我们使用了多语言编码器 umt5-xxl、Mul-OpenCLIP 和 HunyuanDiT_CLIP。 我们在蒸馏练习中接纳了反向去噪处理。

1. import os
2. import torch
3. import torch.nn as nn
5. from typing import Any, Callable, Dict, List, Optional, Union
6. import inspect
7. from diffusers.models.transformers import SD3Transformer2DModel
8. from diffusers.image_processor import VaeImageProcessor
9. from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
10. from diffusers import AutoencoderKL
11. from tqdm import tqdm
12. from PIL import Image
14. from transformers import T5Tokenizer,T5EncoderModel,BertModel, BertTokenizer
15. import open_clip
18. class MLP(nn.Module):
19.     def __init__(self, in_dim=1024, out_dim=2048, hidden_dim=2048, out_dim1=4096, use_residual=True):
20.         super().__init__()
21.         if use_residual:
22.             assert in_dim == out_dim
23.         self.layernorm = nn.LayerNorm(in_dim)
24.         self.projector = nn.Sequential(
25.             nn.Linear(in_dim, hidden_dim, bias=False),
26.             nn.GELU(),
27.             nn.Linear(hidden_dim, hidden_dim, bias=False),
28.             nn.GELU(),
29.             nn.Linear(hidden_dim, hidden_dim, bias=False),
30.             nn.GELU(),
31.             nn.Linear(hidden_dim, out_dim, bias=False),
32.         )
33.         self.fc = nn.Linear(out_dim, out_dim1)
34.         self.use_residual = use_residual
35.     def forward(self, x):
36.         residual = x
37.         x = self.layernorm(x)
38.         x = self.projector(x)
39.         x2 = nn.GELU()(x)
40.         x2 = self.fc(x2)
41.         return x2
43. class Transformer(nn.Module):
44.     def __init__(self, d_model,  n_heads, out_dim1, out_dim2,num_layers=1) -> None:
45.         super().__init__()
47.         self.encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_heads, dim_feedforward=2048, batch_first=True)
48.         self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
49.         self.linear1 = nn.Linear(d_model, out_dim1)
50.         self.linear2 = nn.Linear(d_model, out_dim2)
51.    
52.     def forward(self, x):
53.         x = self.transformer_encoder(x)
54.         x1 = self.linear1(x)
55.         x1 = torch.mean(x1,1)
56.         x2 = self.linear2(x)
57.         return x1,x2
60. def image_grid(imgs, rows, cols):
61.     assert len(imgs) == rows*cols
63.     w, h = imgs[0].size
64.     grid = Image.new('RGB', size=(cols*w, rows*h))
65.     grid_w, grid_h = grid.size
67.     for i, img in enumerate(imgs):
68.         grid.paste(img, box=(i%cols*w, i//cols*h))
69.     return grid
70. def retrieve_timesteps(
71.     scheduler,
72.     num_inference_steps: Optional[int] = None,
73.     device: Optional[Union[str, torch.device]] = None,
74.     timesteps: Optional[List[int]] = None,
75.     sigmas: Optional[List[float]] = None,
76.     **kwargs,
77. ):
78.     if timesteps is not None and sigmas is not None:
79.         raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
80.     if timesteps is not None:
81.         accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
82.         if not accepts_timesteps:
83.             raise ValueError(
84.                 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
85.                 f" timestep schedules. Please check whether you are using the correct scheduler."
86.             )
87.         scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
88.         timesteps = scheduler.timesteps
89.         num_inference_steps = len(timesteps)
90.     elif sigmas is not None:
91.         accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
92.         if not accept_sigmas:
93.             raise ValueError(
94.                 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
95.                 f" sigmas schedules. Please check whether you are using the correct scheduler."
96.             )
97.         scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
98.         timesteps = scheduler.timesteps
99.         num_inference_steps = len(timesteps)
100.     else:
101.         scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
102.         timesteps = scheduler.timesteps
103.     return timesteps, num_inference_steps
105. class StableDiffusionTest():
106.     def __init__(self,model_path,text_encoder_path,text_encoder_path1,text_encoder_path2,proj_path,proj_t5_path):
107.         super().__init__()
108.         self.transformer = SD3Transformer2DModel.from_pretrained(model_path, subfolder="transformer",torch_dtype=dtype).to(device)
109.         self.vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(device,dtype=dtype)
110.         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler")
112.         self.vae_scale_factor = (
113.             2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
114.         )
115.         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
116.         self.default_sample_size = (
117.             self.transformer.config.sample_size
118.             if hasattr(self, "transformer") and self.transformer is not None
119.             else 128
120.         )
122.         self.text_encoder_t5 = T5EncoderModel.from_pretrained(text_encoder_path).to(device,dtype=dtype)
123.         self.tokenizer_t5 = T5Tokenizer.from_pretrained(text_encoder_path)
124.         self.text_encoder = BertModel.from_pretrained(f"{text_encoder_path1}/clip_text_encoder", False, revision=None).to(device,dtype=dtype)
125.         self.tokenizer = BertTokenizer.from_pretrained(f"{text_encoder_path1}/tokenizer")
127.         self.text_encoder2, _, _ = open_clip.create_model_and_transforms('xlm-roberta-large-ViT-H-14', pretrained=text_encoder_path2)
128.         self.tokenizer2 = open_clip.get_tokenizer('xlm-roberta-large-ViT-H-14')
129.         self.text_encoder2.text.output_tokens = True
130.         self.text_encoder2 = self.text_encoder2.to(device,dtype=dtype)
132.         self.proj = MLP(2048, 2048, 2048, 4096, use_residual=False).to(device,dtype=dtype)
133.         self.proj.load_state_dict(torch.load(proj_path, map_location="cpu"))
134.         self.proj_t5 = Transformer(d_model=4096, n_heads=8, out_dim1=2048, out_dim2=4096).to(device,dtype=dtype)
135.         self.proj_t5.load_state_dict(torch.load(proj_t5_path, map_location="cpu"))
137.     def encode_prompt(self, prompt, device, do_classifier_free_guidance=True, negative_prompt=None):
138.         batch_size = len(prompt) if isinstance(prompt, list) else 1
139.         text_input_ids_t5 = self.tokenizer_t5(
140.             prompt,
141.             padding="max_length",
142.             max_length=77,
143.             truncation=True,
144.             add_special_tokens=False,
145.             return_tensors="pt",
146.         ).input_ids.to(device)
148.         text_embeddings = self.text_encoder_t5(text_input_ids_t5)
149.         text_inputs = self.tokenizer(
150.             prompt,
151.             padding="max_length",
152.             max_length=77,
153.             truncation=True,
154.             return_tensors="pt",
155.         )
156.         input_ids = text_inputs.input_ids.to(device)
157.         attention_mask = text_inputs.attention_mask.to(device)
158.         encoder_hidden_states  = self.text_encoder(input_ids,attention_mask=attention_mask)[0]
159.         text_input_ids = self.tokenizer2(prompt).to(device)
160.         _,encoder_hidden_states2  = self.text_encoder2.encode_text(text_input_ids)
161.         encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states2], dim=-1)
163.         encoder_hidden_states_t5 = text_embeddings[0]
164.         encoder_hidden_states = self.proj(encoder_hidden_states)
166.         add_text_embeds,encoder_hidden_states_t5 = self.proj_t5(encoder_hidden_states_t5.half())
167.         prompt_embeds = torch.cat([encoder_hidden_states, encoder_hidden_states_t5], dim=-2)
169.         # get unconditional embeddings for classifier free guidance
170.         if do_classifier_free_guidance:
171.             if negative_prompt is None:
172.                 uncond_tokens = [""] * batch_size
173.             else:
174.                 uncond_tokens = negative_prompt
175.             text_input_ids_t5 = self.tokenizer_t5(
176.                 uncond_tokens,
177.                 padding="max_length",
178.                 max_length=77,
179.                 truncation=True,
180.                 add_special_tokens=False,
181.                 return_tensors="pt",
182.             ).input_ids.to(device)
184.             text_embeddings = self.text_encoder_t5(text_input_ids_t5)
185.             text_inputs = self.tokenizer(
186.                 uncond_tokens,
187.                 padding="max_length",
188.                 max_length=77,
189.                 truncation=True,
190.                 return_tensors="pt",
191.             )
192.             input_ids = text_inputs.input_ids.to(device)
193.             attention_mask = text_inputs.attention_mask.to(device)
194.             encoder_hidden_states  = self.text_encoder(input_ids,attention_mask=attention_mask)[0]
196.             text_input_ids = self.tokenizer2(uncond_tokens).to(device)
197.             _,encoder_hidden_states2  = self.text_encoder2.encode_text(text_input_ids)
198.             encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states2], dim=-1)
200.             encoder_hidden_states_t5 = text_embeddings[0]
201.             encoder_hidden_states_uncond = self.proj(encoder_hidden_states)
203.             add_text_embeds_uncond,encoder_hidden_states_t5_uncond = self.proj_t5(encoder_hidden_states_t5.half())
204.             prompt_embeds_uncond = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_t5_uncond], dim=-2)
206.             prompt_embeds = torch.cat([prompt_embeds_uncond, prompt_embeds], dim=0)
207.             pooled_prompt_embeds = torch.cat([add_text_embeds_uncond, add_text_embeds], dim=0)
209.         return prompt_embeds,pooled_prompt_embeds
212.     def prepare_latents(
213.         self,
214.         batch_size,
215.         num_channels_latents,
216.         height,
217.         width,
218.         dtype,
219.         device,
220.         generator,
221.         latents=None,
222.     ):
223.         if latents is not None:
224.             return latents.to(device=device, dtype=dtype)
226.         shape = (
227.             batch_size,
228.             num_channels_latents,
229.             int(height) // self.vae_scale_factor,
230.             int(width) // self.vae_scale_factor,
231.         )
233.         if isinstance(generator, list) and len(generator) != batch_size:
234.             raise ValueError(
235.                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
236.                 f" size of {batch_size}. Make sure the batch size matches the length of the generators."
237.             )
239.         latents = torch.randn(shape, generator=generator, dtype=dtype).to(device)
241.         return latents
243.     @property
244.     def guidance_scale(self):
245.         return self._guidance_scale
247.     @property
248.     def clip_skip(self):
249.         return self._clip_skip
251.     # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
252.     # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
253.     # corresponds to doing no classifier free guidance.
254.     @property
255.     def do_classifier_free_guidance(self):
256.         return self._guidance_scale > 1
258.     @property
259.     def joint_attention_kwargs(self):
260.         return self._joint_attention_kwargs
262.     @property
263.     def num_timesteps(self):
264.         return self._num_timesteps
266.     @property
267.     def interrupt(self):
268.         return self._interrupt
270.     @torch.no_grad()
271.     def __call__(
272.         self,
273.         prompt: Union[str, List[str]] = None,
274.         prompt_2: Optional[Union[str, List[str]]] = None,
275.         prompt_3: Optional[Union[str, List[str]]] = None,
276.         height: Optional[int] = None,
277.         width: Optional[int] = None,
278.         num_inference_steps: int = 28,
279.         timesteps: List[int] = None,
280.         guidance_scale: float = 7.0,
281.         negative_prompt: Optional[Union[str, List[str]]] = None,
282.         negative_prompt_2: Optional[Union[str, List[str]]] = None,
283.         negative_prompt_3: Optional[Union[str, List[str]]] = None,
284.         num_images_per_prompt: Optional[int] = 1,
285.         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
286.         latents: Optional[torch.FloatTensor] = None,
287.         prompt_embeds: Optional[torch.FloatTensor] = None,
288.         negative_prompt_embeds: Optional[torch.FloatTensor] = None,
289.         pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
290.         negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
291.         output_type: Optional[str] = "pil",
292.         return_dict: bool = True,
293.         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
294.         clip_skip: Optional[int] = None,
295.         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
296.         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
297.     ):
298.         height = height or self.default_sample_size * self.vae_scale_factor
299.         width = width or self.default_sample_size * self.vae_scale_factor
301.         self._guidance_scale = guidance_scale
302.         self._clip_skip = clip_skip
303.         self._joint_attention_kwargs = joint_attention_kwargs
304.         self._interrupt = False
306.         if prompt is not None and isinstance(prompt, str):
307.             batch_size = 1
308.         elif prompt is not None and isinstance(prompt, list):
309.             batch_size = len(prompt)
310.         else:
311.             batch_size = prompt_embeds.shape[0]
314.         prompt_embeds,pooled_prompt_embeds = self.encode_prompt(prompt, device)
316.         timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
317.         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
318.         self._num_timesteps = len(timesteps)
320.         num_channels_latents = self.transformer.config.in_channels
321.         latents = self.prepare_latents(
322.             batch_size * num_images_per_prompt,
323.             num_channels_latents,
324.             height,
325.             width,
326.             prompt_embeds.dtype,
327.             device,
328.             generator,
329.             latents,
330.         )
332.         for i, t in tqdm(enumerate(timesteps)):
333.             if self.interrupt:
334.                 continue
335.             latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
336.             timestep = t.expand(latent_model_input.shape[0]).to(dtype=dtype)
338.             noise_pred = self.transformer(
339.                 hidden_states=latent_model_input,
340.                 timestep=timestep,
341.                 encoder_hidden_states=prompt_embeds.to(dtype=self.transformer.dtype),
342.                 pooled_projections=pooled_prompt_embeds.to(dtype=self.transformer.dtype),
343.                 joint_attention_kwargs=self.joint_attention_kwargs,
344.                 return_dict=False,
345.             )[0]
347.             if self.do_classifier_free_guidance:
348.                 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
349.                 noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
351.             latents_dtype = latents.dtype
352.             latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
354.             if latents.dtype != latents_dtype:
355.                 if torch.backends.mps.is_available():
356.                     latents = latents.to(latents_dtype)
358.             if callback_on_step_end is not None:
359.                 callback_kwargs = {}
360.                 for k in callback_on_step_end_tensor_inputs:
361.                     callback_kwargs[k] = locals()[k]
362.                 callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
364.                 latents = callback_outputs.pop("latents", latents)
365.                 prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
366.                 negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
367.                 negative_pooled_prompt_embeds = callback_outputs.pop(
368.                     "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
369.                 )
371.         if output_type == "latent":
372.             image = latents
373.         else:
374.             latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
375.             image = self.vae.decode(latents, return_dict=False)[0]
376.             image = self.image_processor.postprocess(image, output_type=output_type)
378.         return image
381. if __name__ == '__main__':
382.     device = "cuda"
383.     dtype = torch.float16
385.     text_encoder_path = 'google/umt5-xxl'
386.     text_encoder_path1 = "Tencent-Hunyuan/HunyuanDiT/t2i"
387.     text_encoder_path2 = 'laion/CLIP-ViT-H-14-frozen-xlm-roberta-large-laion5B-s13B-b90k/open_clip_pytorch_model.bin'
389.     model_path = "stabilityai/stable-diffusion-3-medium-diffusers"
390.     proj_path =  "OPPOer/MultilingualSD3-adapter/pytorch_model.bin"
391.     proj_t5_path =  "OPPOer/MultilingualSD3-adapter/pytorch_model_t5.bin"
393.     sdt = StableDiffusionTest(model_path,text_encoder_path,text_encoder_path1,text_encoder_path2,proj_path,proj_t5_path)
395.     batch=2
396.     height = 1024
397.     width = 1024      
398.     while True:
399.         raw_text = input("\nPlease Input Query (stop to exit) >>> ")
400.         if not raw_text:
401.             print('Query should not be empty!')
402.             continue
403.         if raw_text == "stop":
404.             break
405.         images = sdt([raw_text]*batch,height=height,width=width)
406.         grid = image_grid(images, rows=1, cols=batch)
407.         grid.save("MultilingualSD3.png")

复制代码

ChineseFLUX.1-adapter

使用了多语言编码器 byt5-xxl，在顺应过程中使用的西席模型是 FLUX.1-schnell。 我们接纳了反向去噪过程进行蒸馏练习。 理论上，该适配器可应用于任何 FLUX.1 系列模型。 我们在此提供以下应用示例。

1. from diffusers import FluxPipeline, AutoencoderKL
2. from diffusers.image_processor import VaeImageProcessor
3. from transformers import T5ForConditionalGeneration,AutoTokenizer
4. import torch
5. import torch.nn as nn
8. class 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
29. dtype = torch.bfloat16
30. device = "cuda"
31. ckpt_id = "black-forest-labs/FLUX.1-schnell"
32. text_encoder_ckpt_id = 'google/byt5-xxl'
33. proj_t5 = MLP(in_dim=4672, out_dim=4096, hidden_dim=4096, out_dim1=768).to(device=device,dtype=dtype)
34. text_encoder_t5 = T5ForConditionalGeneration.from_pretrained(text_encoder_ckpt_id).get_encoder().to(device=device,dtype=dtype)
35. tokenizer_t5 = AutoTokenizer.from_pretrained(text_encoder_ckpt_id)
38. proj_t5_save_path = f"diffusion_pytorch_model.bin"
39. state_dict = torch.load(proj_t5_save_path, map_location="cpu")
40. state_dict_new = {}
41. for k,v in state_dict.items():
42.     k_new = k.replace("module.","")
43.     state_dict_new[k_new] = v
45. proj_t5.load_state_dict(state_dict_new)
47. pipeline = 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)
53. vae = AutoencoderKL.from_pretrained(
54.     ckpt_id,
55.     subfolder="vae",
56.     torch_dtype=torch.bfloat16
57. ).to(device)
58. vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
59. image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
61. while 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
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
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")

复制代码

MultilingualOpenFLUX.1

OpenFLUX.1
是对 FLUX.1-schnell 模型的微调，已对其进行了蒸馏练习。 请务必更新以下代码中下载的 fast-lora.safetensors 的路径。

1. from diffusers import FluxPipeline, AutoencoderKL
2. from diffusers.image_processor import VaeImageProcessor
3. from transformers import T5ForConditionalGeneration,AutoTokenizer
4. import torch
5. import torch.nn as nn
8. class 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
29. dtype = torch.bfloat16
30. device = "cuda"
31. ckpt_id = "ostris/OpenFLUX.1"
32. text_encoder_ckpt_id = 'google/byt5-xxl'
33. proj_t5 = MLP(in_dim=4672, out_dim=4096, hidden_dim=4096, out_dim1=768).to(device=device,dtype=dtype)
34. text_encoder_t5 = T5ForConditionalGeneration.from_pretrained(text_encoder_ckpt_id).get_encoder().to(device=device,dtype=dtype)
35. tokenizer_t5 = AutoTokenizer.from_pretrained(text_encoder_ckpt_id)
38. proj_t5_save_path = f"diffusion_pytorch_model.bin"
39. state_dict = torch.load(proj_t5_save_path, map_location="cpu")
40. state_dict_new = {}
41. for k,v in state_dict.items():
42.     k_new = k.replace("module.","")
43.     state_dict_new[k_new] = v
45. proj_t5.load_state_dict(state_dict_new)
47. pipeline = 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. pipeline.load_lora_weights("ostris/OpenFLUX.1/openflux1-v0.1.0-fast-lora.safetensors")
54. vae = AutoencoderKL.from_pretrained(
55.     ckpt_id,
56.     subfolder="vae",
57.     torch_dtype=torch.bfloat16
58. ).to(device)
59. vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
60. image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
62. while 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
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
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")

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。