|
--- |
|
license: apache-2.0 |
|
language: |
|
- zh |
|
- en |
|
- fr |
|
- de |
|
- ja |
|
- kg |
|
base_model: |
|
- stabilityai/stable-diffusion-xl-base-1.0 |
|
pipeline_tag: text-to-image |
|
--- |
|
|
|
|
|
![FLUX.1 [schnell] Grid](./PEA-Diffusion.png) |
|
|
|
|
|
Text-to-image diffusion models are well-known for their ability to generate realistic images based on textual prompts. However, the existing works have predominantly focused on English, lacking support for non-English text-to-image models. The most commonly used translation methods cannot solve the generation problem related to language culture, while training from scratch on a specific language dataset is prohibitively expensive. In this paper, we are inspired to propose a simple plug-and-play language transfer method based on knowledge distillation. All we need to do is train a lightweight MLP-like parameter-efficient adapter (PEA) with only 6M parameters under teacher knowledge distillation along with a small parallel data corpus. We are surprised to find that freezing the parameters of UNet can still achieve remarkable performance on the language-specific prompt evaluation set, demonstrating that PEA can stimulate the potential generation ability of the original UNet. Additionally, it closely approaches the performance of the English text-to-image model on a general prompt evaluation set. Furthermore, our adapter can be used as a plugin to achieve significant results in downstream tasks in cross-lingual text-to-image generation. |
|
|
|
# Usage |
|
We provide examples of adapters for models such as [SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0), [Playground v2.5](https://huggingface.co/playgroundai/playground-v2.5-1024px-aesthetic), and [stable-cascade](https://huggingface.co/stabilityai/stable-cascade). For SD3, please refer directly to https://huggingface.co/OPPOer/MultilingualSD3-adapter, and for FLUX. 1, please refer to https://huggingface.co/OPPOer/MultilingualFLUX.1-adapter |
|
|
|
|
|
|
|
## `SDXL` |
|
We used the multilingual encoder [Mul-OpenCLIP](https://huggingface.co/laion/CLIP-ViT-H-14-frozen-xlm-roberta-large-laion5B-s13B-b90k). |
|
As mentioned in the article, you can replace the model here with any SDXL derived model, including sampling acceleration, which can also be directly adapted. |
|
|
|
```python |
|
import os |
|
import torch |
|
import torch.nn as nn |
|
|
|
from PIL import Image |
|
from diffusers import AutoencoderKL, StableDiffusionXLPipeline,DPMSolverMultistepScheduler |
|
from diffusers.image_processor import VaeImageProcessor |
|
from diffusers.models.attention_processor import ( |
|
AttnProcessor2_0, |
|
LoRAAttnProcessor2_0, |
|
LoRAXFormersAttnProcessor, |
|
XFormersAttnProcessor, |
|
) |
|
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Union |
|
import open_clip |
|
|
|
|
|
def image_grid(imgs, rows, cols): |
|
assert len(imgs) == rows*cols |
|
|
|
w, h = imgs[0].size |
|
grid = Image.new('RGB', size=(cols*w, rows*h)) |
|
grid_w, grid_h = grid.size |
|
|
|
for i, img in enumerate(imgs): |
|
grid.paste(img, box=(i%cols*w, i//cols*h)) |
|
return grid |
|
|
|
class MLP(nn.Module): |
|
def __init__(self, in_dim, out_dim, hidden_dim,out_dim1, use_residual=True): |
|
super().__init__() |
|
if use_residual: |
|
assert in_dim == out_dim |
|
self.layernorm = nn.LayerNorm(in_dim) |
|
self.fc1 = nn.Linear(in_dim, hidden_dim) |
|
self.fc2 = nn.Linear(hidden_dim, out_dim) |
|
self.fc3 = nn.Linear(out_dim, out_dim1) |
|
self.use_residual = use_residual |
|
self.act_fn = nn.GELU() |
|
|
|
def forward(self, x): |
|
residual = x |
|
x = self.layernorm(x) |
|
x = self.fc1(x) |
|
x = self.act_fn(x) |
|
x = self.fc2(x) |
|
x2 = self.act_fn(x) |
|
x2 = self.fc3(x2) |
|
if self.use_residual: |
|
x = x + residual |
|
x1 = torch.mean(x,1) |
|
return x1,x2 |
|
|
|
|
|
class StableDiffusionTest(): |
|
|
|
def __init__(self, model_id,text_text_encoder_pathpath,proj_path): |
|
super().__init__() |
|
self.text_encoder, _, preprocess = open_clip.create_model_and_transforms('xlm-roberta-large-ViT-H-14', pretrained=text_encoder_path) |
|
self.tokenizer = open_clip.get_tokenizer('xlm-roberta-large-ViT-H-14') |
|
self.text_encoder.text.output_tokens = True |
|
self.text_encoder = self.text_encoder.to(device,dtype=dtype) |
|
|
|
self.vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device) |
|
scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler") |
|
self.pipe = StableDiffusionXLPipeline.from_pretrained(model_id, scheduler=scheduler,torch_dtype=dtype).to(device) |
|
self.image_processor = VaeImageProcessor(vae_scale_factor=self.pipe.vae_scale_factor) |
|
|
|
self.proj = MLP(1024, 1280, 1024,2048, use_residual=False).to(device,dtype=dtype) |
|
self.proj.load_state_dict(torch.load(proj_path, map_location="cpu")) |
|
|
|
|
|
def encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt): |
|
batch_size = len(prompt) if isinstance(prompt, list) else 1 |
|
|
|
text_input_ids = self.tokenizer(prompt).to(device) |
|
_,text_embeddings = self.text_encoder.encode_text(text_input_ids) |
|
|
|
add_text_embeds,text_embeddings_2048 = self.proj(text_embeddings) |
|
|
|
# duplicate text embeddings for each generation per prompt, using mps friendly method |
|
bs_embed, seq_len, _ = text_embeddings.shape |
|
text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1) |
|
text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) |
|
|
|
# get unconditional embeddings for classifier free guidance |
|
if do_classifier_free_guidance: |
|
uncond_tokens: List[str] |
|
if negative_prompt is None: |
|
uncond_tokens = [""] * batch_size |
|
elif type(prompt) is not type(negative_prompt): |
|
raise TypeError( |
|
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" |
|
f" {type(prompt)}." |
|
) |
|
elif isinstance(negative_prompt, str): |
|
uncond_tokens = [negative_prompt] |
|
elif batch_size != len(negative_prompt): |
|
raise ValueError( |
|
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" |
|
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" |
|
" the batch size of `prompt`." |
|
) |
|
else: |
|
uncond_tokens = negative_prompt |
|
|
|
max_length = text_input_ids.shape[-1] |
|
|
|
uncond_input_ids = self.tokenizer(uncond_tokens).to(device) |
|
_,uncond_embeddings = self.text_encoder.encode_text(uncond_input_ids) |
|
|
|
add_text_embeds_uncond,uncond_embeddings_2048 = self.proj(uncond_embeddings) |
|
|
|
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method |
|
seq_len = uncond_embeddings_2048.shape[1] |
|
uncond_embeddings_2048 = uncond_embeddings_2048.repeat(1, num_images_per_prompt, 1) |
|
uncond_embeddings_2048 = uncond_embeddings_2048.view(batch_size * num_images_per_prompt, seq_len, -1) |
|
|
|
text_embeddings_2048 = torch.cat([uncond_embeddings_2048, text_embeddings_2048]) |
|
add_text_embeds = torch.cat([add_text_embeds_uncond, add_text_embeds]) |
|
|
|
return text_embeddings_2048,add_text_embeds |
|
|
|
def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype): |
|
add_time_ids = list(original_size + crops_coords_top_left + target_size) |
|
add_time_ids = torch.tensor([add_time_ids], dtype=dtype) |
|
return add_time_ids |
|
|
|
|
|
@torch.no_grad() |
|
def __call__( |
|
self, |
|
prompt: Union[str, List[str]], |
|
height: Optional[int] = 1024, |
|
width: Optional[int] = 1024, |
|
num_inference_steps: int = 30, |
|
guidance_scale: float = 7.5, |
|
original_size: Optional[Tuple[int, int]] = None, |
|
crops_coords_top_left: Tuple[int, int] = (0, 0), |
|
target_size: Optional[Tuple[int, int]] = None, |
|
cross_attention_kwargs: Optional[Dict[str, Any]] = None, |
|
guidance_rescale: float = 0, |
|
negative_prompt: Optional[Union[str, List[str]]] = None, |
|
num_images_per_prompt: Optional[int] = 1, |
|
eta: float = 0.0, |
|
generator: Optional[torch.Generator] = None, |
|
latents: Optional[torch.FloatTensor] = None, |
|
prompt_embeds: Optional[torch.FloatTensor] = None, |
|
negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
output_type: Optional[str] = "pil", |
|
return_dict: bool = True, |
|
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, |
|
callback_steps: Optional[int] = 1, |
|
**kwargs, |
|
): |
|
# 0. Default height and width to unet |
|
height = height or self.pipe.unet.config.sample_size * self.pipe.vae_scale_factor |
|
width = width or self.pipe.unet.config.sample_size * self.pipe.vae_scale_factor |
|
original_size = original_size or (height, width) |
|
target_size = target_size or (height, width) |
|
|
|
# 1. Check inputs. Raise error if not correct |
|
# self.pipe.check_inputs(prompt, height, width, callback_steps) |
|
|
|
# 2. Define call parameters |
|
batch_size = 1 if isinstance(prompt, str) else len(prompt) |
|
device = self.pipe._execution_device |
|
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) |
|
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` |
|
# corresponds to doing no classifier free guidance. |
|
do_classifier_free_guidance = guidance_scale > 1.0 |
|
|
|
# 3. Encode input prompt |
|
|
|
prompt_embeds,add_text_embeds = self.encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt) |
|
prompt_embeds = prompt_embeds |
|
add_text_embeds = add_text_embeds |
|
|
|
# 4. Prepare timesteps |
|
self.pipe.scheduler.set_timesteps(num_inference_steps, device=device) |
|
timesteps = self.pipe.scheduler.timesteps |
|
|
|
# 5. Prepare latent variables |
|
num_channels_latents = self.pipe.unet.in_channels |
|
latents = self.pipe.prepare_latents( |
|
batch_size * num_images_per_prompt, |
|
num_channels_latents, |
|
height, |
|
width, |
|
prompt_embeds.dtype, |
|
device, |
|
generator, |
|
latents, |
|
) |
|
|
|
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline |
|
extra_step_kwargs = self.pipe.prepare_extra_step_kwargs(generator, eta) |
|
|
|
add_time_ids = self._get_add_time_ids(original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype) |
|
if do_classifier_free_guidance: |
|
add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0) |
|
|
|
add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) |
|
added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} |
|
|
|
# 7. Denoising loop |
|
for i, t in enumerate(self.pipe.progress_bar(timesteps)): |
|
# expand the latents if we are doing classifier free guidance |
|
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents |
|
latent_model_input = self.pipe.scheduler.scale_model_input(latent_model_input, t) |
|
|
|
# predict the noise residual |
|
noise_pred = self.pipe.unet( |
|
latent_model_input, |
|
t, |
|
encoder_hidden_states=prompt_embeds, |
|
cross_attention_kwargs=cross_attention_kwargs, |
|
added_cond_kwargs=added_cond_kwargs, |
|
return_dict=False, |
|
)[0] |
|
|
|
# noise_pred = self.pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample |
|
|
|
# perform guidance |
|
if do_classifier_free_guidance: |
|
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) |
|
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) |
|
|
|
# compute the previous noisy sample x_t -> x_t-1 |
|
# latents = self.pipe.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample |
|
latents = self.pipe.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] |
|
|
|
# call the callback, if provided |
|
if callback is not None and i % callback_steps == 0: |
|
callback(i, t, latents) |
|
|
|
self.vae.to(dtype=torch.float32) |
|
|
|
use_torch_2_0_or_xformers = self.vae.decoder.mid_block.attentions[0].processor in [ |
|
AttnProcessor2_0, |
|
XFormersAttnProcessor, |
|
LoRAXFormersAttnProcessor, |
|
LoRAAttnProcessor2_0, |
|
] |
|
# if xformers or torch_2_0 is used attention block does not need |
|
# to be in float32 which can save lots of memory |
|
if not use_torch_2_0_or_xformers: |
|
self.vae.post_quant_conv.to(latents.dtype) |
|
self.vae.decoder.conv_in.to(latents.dtype) |
|
self.vae.decoder.mid_block.to(latents.dtype) |
|
else: |
|
latents = latents.float() |
|
|
|
# 8. Post-processing |
|
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] |
|
image = self.image_processor.postprocess(image, output_type="np") |
|
|
|
# 10. Convert to PIL |
|
if output_type == "pil": |
|
image = self.pipe.numpy_to_pil(image) |
|
|
|
return image |
|
|
|
|
|
if __name__ == '__main__': |
|
device = "cuda" |
|
dtype = torch.float16 |
|
|
|
text_encoder_path = 'laion/CLIP-ViT-H-14-frozen-xlm-roberta-large-laion5B-s13B-b90k/open_clip_pytorch_model.bin' |
|
model_id = "stablediffusionapi/protovision-xl-v6.6" |
|
proj_path = "OPPOer/PEA-Diffusion/pytorch_model.bin" |
|
|
|
sdt = StableDiffusionTest(model_id,text_encoder_path,proj_path) |
|
|
|
batch=2 |
|
height = 1024 |
|
width = 1024 |
|
while True: |
|
raw_text = input("\nPlease Input Query (stop to exit) >>> ") |
|
if not raw_text: |
|
print('Query should not be empty!') |
|
continue |
|
if raw_text == "stop": |
|
break |
|
images = sdt([raw_text]*batch,height=height,width=width) |
|
grid = image_grid(images, rows=1, cols=batch) |
|
grid.save("SDXL.png") |
|
|
|
``` |
|
|
|
|
|
|
|
|
|
|
|
## `Playground v2.5` |
|
We used the multilingual encoder [Mul-OpenCLIP](https://huggingface.co/laion/CLIP-ViT-H-14-frozen-xlm-roberta-large-laion5B-s13B-b90k) |
|
|
|
```python |
|
import os,sys |
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Union |
|
import sys |
|
import random |
|
from tqdm import tqdm |
|
|
|
import torch |
|
import torch.nn as nn |
|
import numpy as np |
|
|
|
import argparse |
|
from PIL import Image |
|
import json |
|
from diffusers import AutoencoderKL, DiffusionPipeline |
|
from diffusers.image_processor import VaeImageProcessor |
|
from diffusers.models.attention_processor import ( |
|
AttnProcessor2_0, |
|
LoRAAttnProcessor2_0, |
|
LoRAXFormersAttnProcessor, |
|
XFormersAttnProcessor, |
|
) |
|
import open_clip |
|
|
|
|
|
def image_grid(imgs, rows, cols): |
|
assert len(imgs) == rows*cols |
|
|
|
w, h = imgs[0].size |
|
grid = Image.new('RGB', size=(cols*w, rows*h)) |
|
grid_w, grid_h = grid.size |
|
|
|
for i, img in enumerate(imgs): |
|
grid.paste(img, box=(i%cols*w, i//cols*h)) |
|
return grid |
|
|
|
|
|
class MLP(nn.Module): |
|
def __init__(self, in_dim=1024, out_dim=1280, hidden_dim=2048, out_dim1=2048, use_residual=True): |
|
super().__init__() |
|
if use_residual: |
|
assert in_dim == out_dim |
|
self.layernorm = nn.LayerNorm(in_dim) |
|
self.projector = nn.Sequential( |
|
nn.Linear(in_dim, hidden_dim, bias=False), |
|
nn.GELU(), |
|
nn.Linear(hidden_dim, hidden_dim, bias=False), |
|
nn.GELU(), |
|
nn.Linear(hidden_dim, hidden_dim, bias=False), |
|
nn.GELU(), |
|
nn.Linear(hidden_dim, out_dim, bias=False), |
|
) |
|
self.fc = nn.Linear(out_dim, out_dim1) |
|
self.use_residual = use_residual |
|
def forward(self, x): |
|
residual = x |
|
x = self.layernorm(x) |
|
x = self.projector(x) |
|
x2 = nn.GELU()(x) |
|
x2 = self.fc(x2) |
|
if self.use_residual: |
|
x = x + residual |
|
x1 = torch.mean(x,1) |
|
return x1,x2 |
|
|
|
|
|
class StableDiffusionTest(): |
|
def __init__(self, model_id,text_encoder_path,proj_path): |
|
super().__init__() |
|
self.text_encoder, _, preprocess = open_clip.create_model_and_transforms('xlm-roberta-large-ViT-H-14', pretrained=text_encoder_path) |
|
self.tokenizer = open_clip.get_tokenizer('xlm-roberta-large-ViT-H-14') |
|
self.text_encoder.text.output_tokens = True |
|
self.text_encoder = self.text_encoder.to(device,dtype=dtype) |
|
self.vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device) |
|
|
|
self.pipe = DiffusionPipeline.from_pretrained(model_id, subfolder="scheduler", torch_dtype=dtype, variant="fp16").to(device) |
|
self.image_processor = VaeImageProcessor(vae_scale_factor=self.pipe.vae_scale_factor) |
|
|
|
self.proj = MLP(1024, 1280, 2048, 2048, use_residual=False).to(device,dtype=dtype) |
|
self.proj.load_state_dict(torch.load(proj_path, map_location="cpu")) |
|
|
|
def encode_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt): |
|
batch_size = len(prompt) if isinstance(prompt, list) else 1 |
|
text_input_ids = self.tokenizer(prompt).to(device) |
|
_,text_embeddings = self.text_encoder.encode_text(text_input_ids) |
|
add_text_embeds,text_embeddings_2048 = self.proj(text_embeddings) |
|
|
|
bs_embed, seq_len, _ = text_embeddings.shape |
|
text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1) |
|
text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) |
|
|
|
if do_classifier_free_guidance: |
|
uncond_tokens: List[str] |
|
if negative_prompt is None: |
|
uncond_tokens = [""] * batch_size |
|
elif type(prompt) is not type(negative_prompt): |
|
raise TypeError( |
|
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" |
|
f" {type(prompt)}." |
|
) |
|
elif isinstance(negative_prompt, str): |
|
uncond_tokens = [negative_prompt] |
|
elif batch_size != len(negative_prompt): |
|
raise ValueError( |
|
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" |
|
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" |
|
" the batch size of `prompt`." |
|
) |
|
else: |
|
uncond_tokens = negative_prompt |
|
|
|
max_length = text_input_ids.shape[-1] |
|
uncond_input_ids = self.tokenizer(uncond_tokens).to(device) |
|
_,uncond_embeddings = self.text_encoder.encode_text(uncond_input_ids) |
|
add_text_embeds_uncond,uncond_embeddings_2048 = self.proj(uncond_embeddings) |
|
|
|
seq_len = uncond_embeddings_2048.shape[1] |
|
uncond_embeddings_2048 = uncond_embeddings_2048.repeat(1, num_images_per_prompt, 1) |
|
uncond_embeddings_2048 = uncond_embeddings_2048.view(batch_size * num_images_per_prompt, seq_len, -1) |
|
|
|
text_embeddings_2048 = torch.cat([uncond_embeddings_2048, text_embeddings_2048]) |
|
add_text_embeds = torch.cat([add_text_embeds_uncond, add_text_embeds]) |
|
|
|
return text_embeddings_2048,add_text_embeds |
|
|
|
def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype): |
|
add_time_ids = list(original_size + crops_coords_top_left + target_size) |
|
add_time_ids = torch.tensor([add_time_ids], dtype=dtype) |
|
return add_time_ids |
|
|
|
|
|
@torch.no_grad() |
|
def __call__( |
|
self, |
|
prompt: Union[str, List[str]], |
|
height: Optional[int] = 1024, |
|
width: Optional[int] = 1024, |
|
num_inference_steps: int = 50, |
|
guidance_scale: float = 3, |
|
original_size: Optional[Tuple[int, int]] = None, |
|
crops_coords_top_left: Tuple[int, int] = (0, 0), |
|
target_size: Optional[Tuple[int, int]] = None, |
|
cross_attention_kwargs: Optional[Dict[str, Any]] = None, |
|
guidance_rescale: float = 0, |
|
negative_prompt: Optional[Union[str, List[str]]] = None, |
|
num_images_per_prompt: Optional[int] = 1, |
|
eta: float = 0.0, |
|
generator: Optional[torch.Generator] = None, |
|
latents: Optional[torch.FloatTensor] = None, |
|
prompt_embeds: Optional[torch.FloatTensor] = None, |
|
negative_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, |
|
output_type: Optional[str] = "pil", |
|
return_dict: bool = True, |
|
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, |
|
callback_steps: Optional[int] = 1, |
|
**kwargs, |
|
): |
|
height = height or self.pipe.unet.config.sample_size * self.pipe.vae_scale_factor |
|
width = width or self.pipe.unet.config.sample_size * self.pipe.vae_scale_factor |
|
original_size = original_size or (height, width) |
|
target_size = target_size or (height, width) |
|
|
|
batch_size = 1 if isinstance(prompt, str) else len(prompt) |
|
device = self.pipe._execution_device |
|
|
|
do_classifier_free_guidance = guidance_scale > 1.0 |
|
|
|
prompt_embeds,add_text_embeds = self.encode_prompt(prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt) |
|
|
|
self.pipe.scheduler.set_timesteps(num_inference_steps, device=device) |
|
timesteps = self.pipe.scheduler.timesteps |
|
num_channels_latents = self.pipe.unet.in_channels |
|
latents = self.pipe.prepare_latents( |
|
batch_size * num_images_per_prompt, |
|
num_channels_latents, |
|
height, |
|
width, |
|
prompt_embeds.dtype, |
|
device, |
|
generator, |
|
latents, |
|
) |
|
|
|
extra_step_kwargs = self.pipe.prepare_extra_step_kwargs(generator, eta) |
|
|
|
add_time_ids = self._get_add_time_ids(original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype) |
|
if do_classifier_free_guidance: |
|
add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0) |
|
|
|
add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) |
|
added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} |
|
|
|
for i, t in enumerate(self.pipe.progress_bar(timesteps)): |
|
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents |
|
latent_model_input = self.pipe.scheduler.scale_model_input(latent_model_input, t) |
|
|
|
noise_pred = self.pipe.unet( |
|
latent_model_input, |
|
t, |
|
encoder_hidden_states=prompt_embeds, |
|
cross_attention_kwargs=cross_attention_kwargs, |
|
added_cond_kwargs=added_cond_kwargs, |
|
return_dict=False, |
|
)[0] |
|
|
|
if do_classifier_free_guidance: |
|
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) |
|
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) |
|
|
|
latents = self.pipe.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] |
|
|
|
if callback is not None and i % callback_steps == 0: |
|
callback(i, t, latents) |
|
|
|
self.vae.to(dtype=torch.float32) |
|
|
|
use_torch_2_0_or_xformers = self.vae.decoder.mid_block.attentions[0].processor in [ |
|
AttnProcessor2_0, |
|
XFormersAttnProcessor, |
|
LoRAXFormersAttnProcessor, |
|
LoRAAttnProcessor2_0, |
|
] |
|
|
|
if not use_torch_2_0_or_xformers: |
|
self.vae.post_quant_conv.to(latents.dtype) |
|
self.vae.decoder.conv_in.to(latents.dtype) |
|
self.vae.decoder.mid_block.to(latents.dtype) |
|
else: |
|
latents = latents.float() |
|
|
|
has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None |
|
has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None |
|
if has_latents_mean and has_latents_std: |
|
latents_mean = ( |
|
torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype) |
|
) |
|
latents_std = ( |
|
torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype) |
|
) |
|
latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean |
|
else: |
|
latents = latents / self.vae.config.scaling_factor |
|
|
|
image = self.vae.decode(latents, return_dict=False)[0] |
|
image = self.image_processor.postprocess(image, output_type="np") |
|
|
|
if output_type == "pil": |
|
image = self.pipe.numpy_to_pil(image) |
|
|
|
return image |
|
|
|
|
|
if __name__ == '__main__': |
|
device = "cuda" |
|
dtype = torch.float16 |
|
|
|
model_id = "playgroundai/playground-v2.5-1024px-aesthetic" |
|
text_encoder_path = 'laion/CLIP-ViT-H-14-frozen-xlm-roberta-large-laion5B-s13B-b90k/open_clip_pytorch_model.bin' |
|
proj_path = "OPPOer/PEA-Diffusion/pytorch_model_pg.bin" |
|
|
|
sdt = StableDiffusionTest(model_id,text_encoder_path,proj_path) |
|
|
|
batch=2 |
|
height = 1024 |
|
width = 1024 |
|
|
|
while True: |
|
raw_text = input("\nPlease Input Query (stop to exit) >>> ") |
|
if not raw_text: |
|
print('Query should not be empty!') |
|
continue |
|
if raw_text == "stop": |
|
break |
|
images = sdt([raw_text]*batch,height=height,width=width) |
|
grid = image_grid(images, rows=1, cols=batch) |
|
grid.save("PG.png") |
|
|
|
|
|
``` |
|
To learn more check out the [diffusers](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux) documentation |
|
|
|
|
|
|
|
## `stable-cascade` |
|
comig soon |
|
|
|
|
|
|
|
# License |
|
The adapter itself is Apache License 2.0, but it must follow the license of the main model. |
|
|
|
|
|
# Citation |
|
``` |
|
@misc{ma2023peadiffusion, |
|
title={PEA-Diffusion: Parameter-Efficient Adapter with Knowledge Distillation in non-English Text-to-Image Generation}, |
|
author={Jian Ma and Chen Chen and Qingsong Xie and Haonan Lu}, |
|
year={2023}, |
|
eprint={2311.17086}, |
|
archivePrefix={arXiv}, |
|
primaryClass={cs.CV} |
|
} |
|
``` |
