Diffusers 文件
aMUSEd
並獲得增強的文件體驗
開始使用
aMUSEd
aMUSEd 由 Suraj Patil、William Berman、Robin Rombach 和 Patrick von Platen 在 aMUSEd: An Open MUSE Reproduction 中提出。
Amused 是一種輕量級文字到影像模型,基於 MUSE 架構。Amused 在需要輕量級快速模型的應用中特別有用,例如一次性快速生成多張影像。
Amused 是一種基於 vqvae 令牌的 transformer,與許多擴散模型相比,它能以更少的前向傳播生成影像。與 muse 不同,它使用較小的文字編碼器 CLIP-L/14 而不是 t5-xxl。由於其引數數量少且前向傳播生成過程少,amused 可以快速生成許多影像。這種優勢在大批次處理時尤為明顯。
論文摘要如下:
我們推出了 aMUSEd,這是一款基於 MUSE 的開源輕量級蒙版影像模型 (MIM),用於文字到影像生成。aMUSEd 擁有 MUSE 10% 的引數,專注於快速影像生成。我們認為與潛在擴散(主流的文字到影像生成方法)相比,MIM 尚未得到充分探索。與潛在擴散相比,MIM 需要更少的推理步驟,並且更具可解釋性。此外,MIM 僅需一張影像即可進行微調以學習其他風格。我們希望透過展示 MIM 在大規模文字到影像生成中的有效性併發布可復現的訓練程式碼來鼓勵對 MIM 進行進一步探索。我們還發布了兩種模型的檢查點,它們直接生成 256x256 和 512x512 解析度的影像。
模型 | 引數 |
---|---|
amused-256 | 603M |
amused-512 | 608M |
AmusedPipeline
class diffusers.AmusedPipeline
< 原始碼 >( vqvae: VQModel tokenizer: CLIPTokenizer text_encoder: CLIPTextModelWithProjection transformer: UVit2DModel scheduler: AmusedScheduler )
__call__
< 原始碼 >( prompt: typing.Union[str, typing.List[str], NoneType] = None height: typing.Optional[int] = None width: typing.Optional[int] = None num_inference_steps: int = 12 guidance_scale: float = 10.0 negative_prompt: typing.Union[str, typing.List[str], NoneType] = None num_images_per_prompt: typing.Optional[int] = 1 generator: typing.Optional[torch._C.Generator] = None latents: typing.Optional[torch.IntTensor] = None prompt_embeds: typing.Optional[torch.Tensor] = None encoder_hidden_states: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None negative_encoder_hidden_states: typing.Optional[torch.Tensor] = None output_type = 'pil' return_dict: bool = True callback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = None callback_steps: int = 1 cross_attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None micro_conditioning_aesthetic_score: int = 6 micro_conditioning_crop_coord: typing.Tuple[int, int] = (0, 0) temperature: typing.Union[int, typing.Tuple[int, int], typing.List[int]] = (2, 0) ) → ImagePipelineOutput 或 tuple
引數
- prompt (
str
或List[str]
,可選) — 用於引導影像生成的提示詞。如果未定義,則需要傳遞prompt_embeds
。 - height (
int
,可選,預設為self.transformer.config.sample_size * self.vae_scale_factor
) — 生成影像的高度(以畫素為單位)。 - width (
int
,可選,預設為self.unet.config.sample_size * self.vae_scale_factor
) — 生成影像的寬度(以畫素為單位)。 - num_inference_steps (
int
,可選,預設為 16) — 去噪步數。去噪步數越多通常會生成更高質量的影像,但推理速度會變慢。 - guidance_scale (
float
,可選,預設為 10.0) — 較高的指導比例值會促使模型生成與文字prompt
緊密相關的影像,但會以較低的影像質量為代價。當guidance_scale > 1
時啟用指導比例。 - negative_prompt (
str
或List[str]
,可選) — 引導影像生成中不包含的內容的提示詞。如果未定義,則需要傳遞negative_prompt_embeds
。當不使用指導時 (guidance_scale < 1
) 忽略。 - num_images_per_prompt (
int
,可選,預設為 1) — 每個提示生成的影像數量。 - generator (
torch.Generator
,可選) — 用於使生成具有確定性的torch.Generator
。 - latents (
torch.IntTensor
,可選) — 預生成的令牌,表示self.vqvae
中的潛在向量,用作影像生成的輸入。如果未提供,起始潛在向量將完全被遮罩。 - prompt_embeds (
torch.Tensor
,可選) — 預生成的文字嵌入。可用於輕鬆調整文字輸入(提示權重)。如果未提供,文字嵌入將從prompt
輸入引數生成。來自池化和投影的最終隱藏狀態的單個向量。 - encoder_hidden_states (
torch.Tensor
,可選) — 預生成的來自文字編碼器的倒數第二個隱藏狀態,提供額外的文字條件。 - negative_prompt_embeds (
torch.Tensor
,可選) — 預生成的負面文字嵌入。可用於輕鬆調整文字輸入(提示權重)。如果未提供,negative_prompt_embeds
將從negative_prompt
輸入引數生成。 - negative_encoder_hidden_states (
torch.Tensor
,可選) — 與積極提示詞的encoder_hidden_states
類似。 - output_type (
str
,可選,預設為"pil"
) — 生成影像的輸出格式。在PIL.Image
或np.array
之間選擇。 - return_dict (
bool
,可選,預設為True
) — 是否返回 StableDiffusionPipelineOutput 而不是普通的元組。 - callback (
Callable
,可選) — 在推理過程中每callback_steps
步呼叫的函式。該函式將使用以下引數呼叫:callback(step: int, timestep: int, latents: torch.Tensor)
。 - callback_steps (
int
,可選,預設為 1) — 呼叫callback
函式的頻率。如果未指定,回撥將在每一步呼叫。 - cross_attention_kwargs (
dict
,可選) — 一個 kwargs 字典,如果指定,將作為引數傳遞給self.processor
中定義的AttentionProcessor
。 - micro_conditioning_aesthetic_score (
int
, 可選, 預設為 6) — 根據 laion 美學分類器設定的目標美學分數。請參閱 https://laion.ai/blog/laion-aesthetics/ 以及 https://huggingface.co/papers/2307.01952 中的微條件化部分。 - micro_conditioning_crop_coord (
Tuple[int]
, 可選, 預設為 (0, 0)) — 目標高度、寬度裁剪座標。請參閱 https://huggingface.co/papers/2307.01952 中的微條件化部分。 - temperature (
Union[int, Tuple[int, int], List[int]]
, 可選, 預設為 (2, 0)) — 配置self.scheduler
上的溫度排程器,請參閱AmusedScheduler#set_timesteps
。
返回
ImagePipelineOutput 或 tuple
如果 return_dict
為 True
,則返回 ImagePipelineOutput,否則返回一個 tuple
,其中第一個元素是生成的影像列表。
用於生成的管道的呼叫函式。
enable_xformers_memory_efficient_attention
< 來源 >( attention_op: typing.Optional[typing.Callable] = None )
引數
- attention_op (
Callable
, 可選) — 覆蓋預設的None
運算子,用作 xFormers 的memory_efficient_attention()
函式的op
引數。
啟用 xFormers 的記憶體高效注意力。啟用此選項後,您應該會觀察到 GPU 記憶體使用量降低,並且推理速度可能會加快。訓練期間的速度提升無法保證。
⚠️ 當記憶體高效注意力和切片注意力同時啟用時,記憶體高效注意力優先。
示例
>>> import torch
>>> from diffusers import DiffusionPipeline
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
>>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
>>> pipe = pipe.to("cuda")
>>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
>>> # Workaround for not accepting attention shape using VAE for Flash Attention
>>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)
class diffusers.AmusedImg2ImgPipeline
< 來源 >( vqvae: VQModel tokenizer: CLIPTokenizer text_encoder: CLIPTextModelWithProjection transformer: UVit2DModel scheduler: AmusedScheduler )
__call__
< 來源 >( prompt: typing.Union[str, typing.List[str], NoneType] = None image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor]] = None strength: float = 0.5 num_inference_steps: int = 12 guidance_scale: float = 10.0 negative_prompt: typing.Union[str, typing.List[str], NoneType] = None num_images_per_prompt: typing.Optional[int] = 1 generator: typing.Optional[torch._C.Generator] = None prompt_embeds: typing.Optional[torch.Tensor] = None encoder_hidden_states: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None negative_encoder_hidden_states: typing.Optional[torch.Tensor] = None output_type = 'pil' return_dict: bool = True callback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = None callback_steps: int = 1 cross_attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None micro_conditioning_aesthetic_score: int = 6 micro_conditioning_crop_coord: typing.Tuple[int, int] = (0, 0) temperature: typing.Union[int, typing.Tuple[int, int], typing.List[int]] = (2, 0) ) → ImagePipelineOutput 或 tuple
引數
- prompt (
str
或List[str]
, 可選) — 用於引導影像生成的提示詞。如果未定義,您需要傳入prompt_embeds
。 - image (
torch.Tensor
,PIL.Image.Image
,np.ndarray
,List[torch.Tensor]
,List[PIL.Image.Image]
, 或List[np.ndarray]
) — 用作起點的影像、Numpy 陣列或表示影像批次的張量。對於 Numpy 陣列和 PyTorch 張量,預期值範圍應在[0, 1]
之間。如果它是張量或張量列表,則預期形狀應為(B, C, H, W)
或(C, H, W)
。如果它是 Numpy 陣列或陣列列表,則預期形狀應為(B, H, W, C)
或(H, W, C)
。它也可以接受影像潛在表示作為image
,但如果直接傳入潛在表示則不會再次編碼。 - strength (
float
, 可選, 預設為 0.5) — 指示參考image
的轉換程度。必須在 0 到 1 之間。image
用作起點,strength
越高,新增的噪聲越多。去噪步驟的數量取決於最初新增的噪聲量。當strength
為 1 時,新增的噪聲最大,去噪過程將執行num_inference_steps
中指定的完整迭代次數。值為 1 基本忽略image
。 - num_inference_steps (
int
, 可選, 預設為 12) — 去噪步驟的數量。更多的去噪步驟通常會生成更高質量的影像,但推理速度會變慢。 - guidance_scale (
float
, 可選, 預設為 10.0) — 更高的引導尺度值鼓勵模型生成與文字prompt
緊密相關的影像,但代價是影像質量會降低。當guidance_scale > 1
時啟用引導尺度。 - negative_prompt (
str
或List[str]
, 可選) — 用於引導影像生成中不包含內容的提示詞。如果未定義,您需要傳入negative_prompt_embeds
。當不使用引導時(guidance_scale < 1
)將被忽略。 - num_images_per_prompt (
int
, 可選, 預設為 1) — 每個提示要生成的影像數量。 - generator (
torch.Generator
, 可選) — 用於使生成確定性的torch.Generator
。 - prompt_embeds (
torch.Tensor
, 可選) — 預生成的文字嵌入。可用於輕鬆調整文字輸入(提示權重)。如果未提供,則從prompt
輸入引數生成文字嵌入。這是來自池化和投影后的最終隱藏狀態的單個向量。 - encoder_hidden_states (
torch.Tensor
, 可選) — 文字編碼器預生成的倒數第二個隱藏狀態,提供額外的文字條件。 - negative_prompt_embeds (
torch.Tensor
, 可選) — 預生成的負面文字嵌入。可用於輕鬆調整文字輸入(提示權重)。如果未提供,則從negative_prompt
輸入引數生成negative_prompt_embeds
。 - negative_encoder_hidden_states (
torch.Tensor
, 可選) — 與正面提示的encoder_hidden_states
類似。 - output_type (
str
, 可選, 預設為"pil"
) — 生成影像的輸出格式。可選擇PIL.Image
或np.array
。 - return_dict (
bool
, 可選, 預設為True
) — 是否返回 StableDiffusionPipelineOutput 而不是普通元組。 - callback (
Callable
, 可選) — 在推理過程中每callback_steps
步呼叫的函式。該函式將以以下引數呼叫:callback(step: int, timestep: int, latents: torch.Tensor)
。 - callback_steps (
int
, 可選, 預設為 1) — 呼叫callback
函式的頻率。如果未指定,回撥將在每一步呼叫。 - cross_attention_kwargs (
dict
, 可選) — 一個 kwargs 字典,如果指定,將作為self.processor
中定義的AttentionProcessor
的op
引數傳入。 - micro_conditioning_aesthetic_score (
int
, 可選, 預設為 6) — 根據 laion 美學分類器設定的目標美學分數。請參閱 https://laion.ai/blog/laion-aesthetics/ 以及 https://huggingface.co/papers/2307.01952 中的微條件化部分。 - micro_conditioning_crop_coord (
Tuple[int]
, 可選, 預設為 (0, 0)) — 目標高度、寬度裁剪座標。請參閱 https://huggingface.co/papers/2307.01952 中的微條件化部分。 - temperature (
Union[int, Tuple[int, int], List[int]]
, 可選, 預設為 (2, 0)) — 配置self.scheduler
上的溫度排程器,請參閱AmusedScheduler#set_timesteps
。
返回
ImagePipelineOutput 或 tuple
如果 return_dict
為 True
,則返回 ImagePipelineOutput,否則返回一個 tuple
,其中第一個元素是生成的影像列表。
用於生成的管道的呼叫函式。
示例
>>> import torch
>>> from diffusers import AmusedImg2ImgPipeline
>>> from diffusers.utils import load_image
>>> pipe = AmusedImg2ImgPipeline.from_pretrained(
... "amused/amused-512", variant="fp16", torch_dtype=torch.float16
... )
>>> pipe = pipe.to("cuda")
>>> prompt = "winter mountains"
>>> input_image = (
... load_image(
... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/open_muse/mountains.jpg"
... )
... .resize((512, 512))
... .convert("RGB")
... )
>>> image = pipe(prompt, input_image).images[0]
enable_xformers_memory_efficient_attention
< 來源 >( attention_op: typing.Optional[typing.Callable] = None )
引數
- attention_op (
Callable
, 可選) — 覆蓋預設的None
運算子,用作 xFormers 的memory_efficient_attention()
函式的op
引數。
啟用 xFormers 的記憶體高效注意力。啟用此選項後,您應該會觀察到 GPU 記憶體使用量降低,並且推理速度可能會加快。訓練期間的速度提升無法保證。
⚠️ 當記憶體高效注意力和切片注意力同時啟用時,記憶體高效注意力優先。
示例
>>> import torch
>>> from diffusers import DiffusionPipeline
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
>>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
>>> pipe = pipe.to("cuda")
>>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
>>> # Workaround for not accepting attention shape using VAE for Flash Attention
>>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)
class diffusers.AmusedInpaintPipeline
< 來源 >( vqvae: VQModel tokenizer: CLIPTokenizer text_encoder: CLIPTextModelWithProjection transformer: UVit2DModel scheduler: AmusedScheduler )
__call__
< 來源 >( prompt: typing.Union[str, typing.List[str], NoneType] = None image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor]] = None mask_image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor]] = None strength: float = 1.0 num_inference_steps: int = 12 guidance_scale: float = 10.0 negative_prompt: typing.Union[str, typing.List[str], NoneType] = None num_images_per_prompt: typing.Optional[int] = 1 generator: typing.Optional[torch._C.Generator] = None prompt_embeds: typing.Optional[torch.Tensor] = None encoder_hidden_states: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None negative_encoder_hidden_states: typing.Optional[torch.Tensor] = None output_type = 'pil' return_dict: bool = True callback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = None callback_steps: int = 1 cross_attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None micro_conditioning_aesthetic_score: int = 6 micro_conditioning_crop_coord: typing.Tuple[int, int] = (0, 0) temperature: typing.Union[int, typing.Tuple[int, int], typing.List[int]] = (2, 0) ) → ImagePipelineOutput 或 tuple
引數
- prompt (
str
或List[str]
, 可選) — 用於引導影像生成的提示詞。如果未定義,您需要傳入prompt_embeds
。 - image (
torch.Tensor
,PIL.Image.Image
,np.ndarray
,List[torch.Tensor]
,List[PIL.Image.Image]
, 或List[np.ndarray]
) — 用作起點的影像、Numpy 陣列或表示影像批次的張量。對於 Numpy 陣列和 PyTorch 張量,預期值範圍應在[0, 1]
之間。如果它是張量或張量列表,則預期形狀應為(B, C, H, W)
或(C, H, W)
。如果它是 Numpy 陣列或陣列列表,則預期形狀應為(B, H, W, C)
或(H, W, C)
。它也可以接受影像潛在表示作為image
,但如果直接傳入潛在表示則不會再次編碼。 - mask_image (
torch.Tensor
,PIL.Image.Image
,np.ndarray
,List[torch.Tensor]
,List[PIL.Image.Image]
, 或List[np.ndarray]
) — 影像、Numpy 陣列或表示影像批次的張量,用於遮蓋image
。遮罩中的白色畫素將被重新繪製,而黑色畫素將保留。如果mask_image
是 PIL 影像,它在使用前將轉換為單通道(亮度)。如果它是 Numpy 陣列或 PyTorch 張量,它應該包含一個顏色通道(L)而不是 3 個,因此 PyTorch 張量的預期形狀為(B, 1, H, W)
、(B, H, W)
、(1, H, W)
或(H, W)
。對於 Numpy 陣列,預期形狀為(B, H, W, 1)
、(B, H, W)
、(H, W, 1)
或(H, W)
。 - strength (
float
, 可選, 預設為 1.0) — 指示參考image
的轉換程度。必須在 0 到 1 之間。image
用作起點,strength
越高,新增的噪聲越多。去噪步驟的數量取決於最初新增的噪聲量。當strength
為 1 時,新增的噪聲最大,去噪過程將執行num_inference_steps
中指定的完整迭代次數。值為 1 基本忽略image
。 - num_inference_steps (
int
, 可選, 預設為 16) — 去噪步驟的數量。更多的去噪步驟通常會生成更高質量的影像,但推理速度會變慢。 - guidance_scale (
float
, 可選, 預設為 10.0) — 更高的引導尺度值鼓勵模型生成與文字prompt
緊密相關的影像,但代價是影像質量會降低。當guidance_scale > 1
時啟用引導尺度。 - negative_prompt (
str
或List[str]
, 可選) — 用於引導影像生成中不包含內容的提示詞。如果未定義,您需要傳入negative_prompt_embeds
。當不使用引導時(guidance_scale < 1
)將被忽略。 - num_images_per_prompt (
int
, 可選, 預設為 1) — 每個提示要生成的影像數量。 - generator (
torch.Generator
, 可選) — 用於使生成確定性的torch.Generator
。 - prompt_embeds (
torch.Tensor
, 可選) — 預生成的文字嵌入。可用於輕鬆調整文字輸入(提示權重)。如果未提供,則從prompt
輸入引數生成文字嵌入。這是來自池化和投影后的最終隱藏狀態的單個向量。 - encoder_hidden_states (
torch.Tensor
, 可選) — 文字編碼器預生成的倒數第二個隱藏狀態,提供額外的文字條件。 - negative_prompt_embeds (
torch.Tensor
, 可選) — 預生成的負面文字嵌入。可用於輕鬆調整文字輸入(提示權重)。如果未提供,則從negative_prompt
輸入引數生成negative_prompt_embeds
。 - negative_encoder_hidden_states (
torch.Tensor
, 可選) — 類似於用於正面提示的encoder_hidden_states
。 - output_type (
str
, 可選, 預設為"pil"
) — 生成影像的輸出格式。可在PIL.Image
或np.array
之間選擇。 - return_dict (
bool
, 可選, 預設為True
) — 是否返回 StableDiffusionPipelineOutput 而非普通元組。 - callback (
Callable
, 可選) — 在推理過程中每callback_steps
步呼叫的函式。該函式呼叫時帶有以下引數:callback(step: int, timestep: int, latents: torch.Tensor)
。 - callback_steps (
int
, 可選, 預設為 1) — 呼叫callback
函式的頻率。如果未指定,則在每一步都呼叫回撥。 - cross_attention_kwargs (
dict
, 可選) — 如果指定,此 kwargs 字典將作為op
引數傳遞給self.processor
中定義的AttentionProcessor
。 - micro_conditioning_aesthetic_score (
int
, 可選, 預設為 6) — 根據 laion 美學分類器設定的目標美學分數。請參閱 https://laion.ai/blog/laion-aesthetics/ 和 https://huggingface.co/papers/2307.01952 中的微條件部分。 - micro_conditioning_crop_coord (
Tuple[int]
, 可選, 預設為 (0, 0)) — 目標高度、寬度裁剪座標。請參閱 https://huggingface.co/papers/2307.01952 中的微條件部分。 - temperature (
Union[int, Tuple[int, int], List[int]]
, 可選, 預設為 (2, 0)) — 配置self.scheduler
上的溫度排程器,請參閱AmusedScheduler#set_timesteps
。
返回
ImagePipelineOutput 或 tuple
如果 return_dict
為 True
,則返回 ImagePipelineOutput,否則返回一個 tuple
,其中第一個元素是生成的影像列表。
用於生成的管道的呼叫函式。
示例
>>> import torch
>>> from diffusers import AmusedInpaintPipeline
>>> from diffusers.utils import load_image
>>> pipe = AmusedInpaintPipeline.from_pretrained(
... "amused/amused-512", variant="fp16", torch_dtype=torch.float16
... )
>>> pipe = pipe.to("cuda")
>>> prompt = "fall mountains"
>>> input_image = (
... load_image(
... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/open_muse/mountains_1.jpg"
... )
... .resize((512, 512))
... .convert("RGB")
... )
>>> mask = (
... load_image(
... "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/open_muse/mountains_1_mask.png"
... )
... .resize((512, 512))
... .convert("L")
... )
>>> pipe(prompt, input_image, mask).images[0].save("out.png")
enable_xformers_memory_efficient_attention
< 來源 >( attention_op: typing.Optional[typing.Callable] = None )
引數
- attention_op (
Callable
, 可選) — 覆蓋預設的None
運算子,用作 xFormers 的memory_efficient_attention()
函式的op
引數。
啟用 xFormers 的記憶體高效注意力。啟用此選項後,您應該會觀察到 GPU 記憶體使用量降低,並且推理速度可能會加快。訓練期間的速度提升無法保證。
⚠️ 當記憶體高效注意力和切片注意力同時啟用時,記憶體高效注意力優先。
示例
>>> import torch
>>> from diffusers import DiffusionPipeline
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
>>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
>>> pipe = pipe.to("cuda")
>>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
>>> # Workaround for not accepting attention shape using VAE for Flash Attention
>>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)