diffusers 原始碼解析（二十六）
.\diffusers\pipelines\deepfloyd_if\pipeline_if_inpainting_superresolution.py
# 匯入 html 模組，用於處理 HTML 文字
import html
# 匯入 inspect 模組，用於獲取物件的資訊
import inspect
# 匯入 re 模組，用於正規表示式匹配
import re
# 匯入 urllib.parse 模組並重新命名為 ul，用於處理 URL 編碼
import urllib.parse as ul
# 從 typing 模組匯入型別提示相關的類
from typing import Any, Callable, Dict, List, Optional, Union
# 匯入 numpy 庫並重新命名為 np，用於陣列和數學計算
import numpy as np
# 匯入 PIL.Image 模組，用於影像處理
import PIL.Image
# 匯入 torch 庫，用於深度學習操作
import torch
# 從 torch.nn.functional 匯入 F，用於神經網路的功能操作
import torch.nn.functional as F
# 從 transformers 庫匯入 CLIPImageProcessor, T5EncoderModel 和 T5Tokenizer，用於處理影像和文字
from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer
# 從本地模組匯入 StableDiffusionLoraLoaderMixin，用於載入穩定擴散模型
from ...loaders import StableDiffusionLoraLoaderMixin
# 從本地模組匯入 UNet2DConditionModel，用於2D條件模型
from ...models import UNet2DConditionModel
# 從本地模組匯入 DDPMScheduler，用於擴散排程
from ...schedulers import DDPMScheduler
# 從本地模組匯入多個實用工具函式
from ...utils import (
BACKENDS_MAPPING, # 後端對映
PIL_INTERPOLATION, # PIL 插值方式
is_bs4_available, # 檢查 BeautifulSoup 是否可用
is_ftfy_available, # 檢查 ftfy 是否可用
logging, # 日誌記錄模組
replace_example_docstring, # 替換示例文件字串的函式
)
# 從本地工具模組匯入 randn_tensor 函式，用於生成隨機張量
from ...utils.torch_utils import randn_tensor
# 從本地模組匯入 DiffusionPipeline，用於處理擴散管道
from ..pipeline_utils import DiffusionPipeline
# 從本地模組匯入 IFPipelineOutput，用於擴散管道的輸出
from .pipeline_output import IFPipelineOutput
# 從本地模組匯入 IFSafetyChecker，用於安全檢查
from .safety_checker import IFSafetyChecker
# 從本地模組匯入 IFWatermarker，用於新增水印
from .watermark import IFWatermarker
# 如果 BeautifulSoup 可用，則匯入 BeautifulSoup 類
if is_bs4_available():
from bs4 import BeautifulSoup
# 如果 ftfy 可用，則匯入 ftfy 模組
if is_ftfy_available():
import ftfy
# 建立一個 logger 例項用於記錄日誌，命名為當前模組名
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.resize 複製的函式
def resize(images: PIL.Image.Image, img_size: int) -> PIL.Image.Image:
# 獲取輸入影像的寬度和高度
w, h = images.size
# 計算寬高比
coef = w / h
# 初始化寬高為目標尺寸
w, h = img_size, img_size
# 如果寬高比大於等於 1，則按比例調整寬度
if coef >= 1:
w = int(round(img_size / 8 * coef) * 8) # 調整寬度為 8 的倍數
else:
# 如果寬高比小於 1，則按比例調整高度
h = int(round(img_size / 8 / coef) * 8) # 調整高度為 8 的倍數
# 調整影像大小，使用雙三次插值法
images = images.resize((w, h), resample=PIL_INTERPOLATION["bicubic"], reducing_gap=None)
# 返回撥整後的影像
return images
# 示例文件字串
EXAMPLE_DOC_STRING = """
# 示例用法
Examples:
```py
# 匯入必要的庫和模組
>>> from diffusers import IFInpaintingPipeline, IFInpaintingSuperResolutionPipeline, DiffusionPipeline
>>> from diffusers.utils import pt_to_pil
>>> import torch
>>> from PIL import Image
>>> import requests
>>> from io import BytesIO
# 定義原始影像的 URL
>>> url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/person.png"
# 傳送 GET 請求以獲取影像
>>> response = requests.get(url)
# 將響應內容轉換為 RGB 格式的影像
>>> original_image = Image.open(BytesIO(response.content)).convert("RGB")
# 將原始影像賦值給原始影像變數
>>> original_image = original_image
# 定義掩膜影像的 URL
>>> url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/glasses_mask.png"
# 傳送 GET 請求以獲取掩膜影像
>>> response = requests.get(url)
# 將響應內容轉換為影像
>>> mask_image = Image.open(BytesIO(response.content))
# 將掩膜影像賦值給掩膜影像變數
>>> mask_image = mask_image
# 從預訓練模型載入影像修復管道
>>> pipe = IFInpaintingPipeline.from_pretrained(
... "DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16
... )
# 啟用模型 CPU 解除安裝功能以節省記憶體
>>> pipe.enable_model_cpu_offload()
# 定義提示文字
>>> prompt = "blue sunglasses"
# 編碼提示文字為嵌入向量，包括正面和負面提示
>>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
# 使用原始影像、掩膜影像和提示嵌入生成影像
>>> image = pipe(
... image=original_image,
... mask_image=mask_image,
... prompt_embeds=prompt_embeds,
... negative_prompt_embeds=negative_embeds,
... output_type="pt",
... ).images
# 儲存中間生成的影像
>>> pil_image = pt_to_pil(image)
# 將中間影像儲存到檔案
>>> pil_image[0].save("./if_stage_I.png")
# 從預訓練模型載入超解析度管道
>>> super_res_1_pipe = IFInpaintingSuperResolutionPipeline.from_pretrained(
... "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16
... )
# 啟用模型 CPU 解除安裝功能以節省記憶體
>>> super_res_1_pipe.enable_model_cpu_offload()
# 使用超解析度管道生成最終影像
>>> image = super_res_1_pipe(
... image=image,
... mask_image=mask_image,
... original_image=original_image,
... prompt_embeds=prompt_embeds,
... negative_prompt_embeds=negative_embeds,
... ).images
# 將最終影像儲存到檔案
>>> image[0].save("./if_stage_II.png")
```py
"""
# 定義一個名為 IFInpaintingSuperResolutionPipeline 的類，繼承自 DiffusionPipeline 和 StableDiffusionLoraLoaderMixin
class IFInpaintingSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
# 定義一個 tokenizer 屬性，型別為 T5Tokenizer
tokenizer: T5Tokenizer
# 定義一個 text_encoder 屬性，型別為 T5EncoderModel
text_encoder: T5EncoderModel
# 定義一個 unet 屬性，型別為 UNet2DConditionModel
unet: UNet2DConditionModel
# 定義一個排程器 scheduler，型別為 DDPMScheduler
scheduler: DDPMScheduler
# 定義一個影像噪聲排程器 image_noising_scheduler，型別為 DDPMScheduler
image_noising_scheduler: DDPMScheduler
# 可選的特徵提取器，型別為 CLIPImageProcessor
feature_extractor: Optional[CLIPImageProcessor]
# 可選的安全檢查器，型別為 IFSafetyChecker
safety_checker: Optional[IFSafetyChecker]
# 可選的水印處理器，型別為 IFWatermarker
watermarker: Optional[IFWatermarker]
# 定義一個正規表示式，用於匹配不良標點
bad_punct_regex = re.compile(
r"["
+ "#®•©™&@·º½¾¿¡§~" # 包含特定特殊字元
+ r"\)" # 匹配右括號
+ r"\(" # 匹配左括號
+ r"\]" # 匹配右中括號
+ r"\[" # 匹配左中括號
+ r"\}" # 匹配右花括號
+ r"\{" # 匹配左花括號
+ r"\|" # 匹配豎線
+ "\\"
+ r"\/" # 匹配斜槓
+ r"\*" # 匹配星號
+ r"]{1,}" # 至少匹配一個以上的字元
) # noqa
# 定義一個字串，用於表示 CPU 解除安裝順序
model_cpu_offload_seq = "text_encoder->unet"
# 定義一個可選元件列表
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
# 定義一個不參與 CPU 解除安裝的元件列表
_exclude_from_cpu_offload = ["watermarker"]
# 初始化方法，接收多個引數以構建類的例項
def __init__(
self,
# tokenizer 引數，型別為 T5Tokenizer
tokenizer: T5Tokenizer,
# text_encoder 引數，型別為 T5EncoderModel
text_encoder: T5EncoderModel,
# unet 引數，型別為 UNet2DConditionModel
unet: UNet2DConditionModel,
# scheduler 引數，型別為 DDPMScheduler
scheduler: DDPMScheduler,
# image_noising_scheduler 引數，型別為 DDPMScheduler
image_noising_scheduler: DDPMScheduler,
# 可選的安全檢查器引數，型別為 IFSafetyChecker
safety_checker: Optional[IFSafetyChecker],
# 可選的特徵提取器引數，型別為 CLIPImageProcessor
feature_extractor: Optional[CLIPImageProcessor],
# 可選的水印處理器引數，型別為 IFWatermarker
watermarker: Optional[IFWatermarker],
# 指示是否需要安全檢查器的布林值，預設為 True
requires_safety_checker: bool = True,
):
# 呼叫父類的初始化方法
super().__init__()
# 檢查安全檢查器是否為 None 且要求使用安全檢查器
if safety_checker is None and requires_safety_checker:
# 記錄警告，提示使用者已禁用安全檢查器
logger.warning(
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
" that you abide to the conditions of the IF license and do not expose unfiltered"
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
)
# 檢查安全檢查器是否不為 None 且特徵提取器為 None
if safety_checker is not None and feature_extractor is None:
# 丟擲值錯誤，提示必須定義特徵提取器
raise ValueError(
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
)
# 檢查 UNet 配置的輸入通道數是否不等於 6
if unet.config.in_channels != 6:
# 記錄警告，提示使用者載入了不適用於超解析度的檢查點
logger.warning(
"It seems like you have loaded a checkpoint that shall not be used for super resolution from {unet.config._name_or_path} as it accepts {unet.config.in_channels} input channels instead of 6. Please make sure to pass a super resolution checkpoint as the `'unet'`: IFSuperResolutionPipeline.from_pretrained(unet=super_resolution_unet, ...)`."
)
# 註冊各個模組以便後續使用
self.register_modules(
tokenizer=tokenizer,
text_encoder=text_encoder,
unet=unet,
scheduler=scheduler,
image_noising_scheduler=image_noising_scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
watermarker=watermarker,
)
# 將配置註冊到例項中，記錄是否需要安全檢查器
self.register_to_config(requires_safety_checker=requires_safety_checker)
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing 複製
# 定義一個文字預處理的私有方法，接收文字和是否清理標題的標誌
def _text_preprocessing(self, text, clean_caption=False):
# 如果設定了清理標題但 bs4 庫不可用，發出警告並將 clean_caption 設定為 False
if clean_caption and not is_bs4_available():
logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
logger.warning("Setting `clean_caption` to False...")
clean_caption = False
# 如果設定了清理標題但 ftfy 庫不可用，發出警告並將 clean_caption 設定為 False
if clean_caption and not is_ftfy_available():
logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
logger.warning("Setting `clean_caption` to False...")
clean_caption = False
# 如果文字不是元組或列表，則將其包裝成列表
if not isinstance(text, (tuple, list)):
text = [text]
# 定義一個處理單個文字的內部函式
def process(text: str):
# 如果需要清理標題，則呼叫清理標題的方法
if clean_caption:
text = self._clean_caption(text)
text = self._clean_caption(text)
else:
# 否則將文字轉換為小寫並去除首尾空白
text = text.lower().strip()
# 返回處理後的文字
return text
# 對文字列表中的每個元素應用處理函式，並返回處理後的結果列表
return [process(t) for t in text]
# 禁用梯度計算，以節省記憶體和提高效能
@torch.no_grad()
# 定義編碼提示的方法，接收多個引數
def encode_prompt(
self,
prompt: Union[str, List[str]],
do_classifier_free_guidance: bool = True,
num_images_per_prompt: int = 1,
device: Optional[torch.device] = None,
negative_prompt: Optional[Union[str, List[str]]] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
clean_caption: bool = False,
# 定義執行安全檢查的方法，接收影像、裝置和資料型別
def run_safety_checker(self, image, device, dtype):
# 如果存在安全檢查器
if self.safety_checker is not None:
# 使用特徵提取器處理影像，並將結果轉換為指定裝置的張量
safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)
# 執行安全檢查器，檢測影像中的不當內容和水印
image, nsfw_detected, watermark_detected = self.safety_checker(
images=image,
clip_input=safety_checker_input.pixel_values.to(dtype=dtype),
)
else:
# 如果沒有安全檢查器，則將檢測結果設定為 None
nsfw_detected = None
watermark_detected = None
# 返回處理後的影像和檢測結果
return image, nsfw_detected, watermark_detected
# 這裡是從其他模組複製的準備額外步驟引數的方法
# 準備排程器步驟的額外引數，因為並非所有排程器都有相同的簽名
def prepare_extra_step_kwargs(self, generator, eta):
# eta (η) 僅在 DDIMScheduler 中使用，其他排程器將忽略它
# eta 對應於 DDIM 論文中的 η: https://arxiv.org/abs/2010.02502
# 應在 [0, 1] 範圍內
# 檢查排程器的 step 方法引數中是否接受 eta
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 初始化額外步驟引數字典
extra_step_kwargs = {}
# 如果接受 eta，則將 eta 新增到額外引數字典中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 檢查排程器的 step 方法引數中是否接受 generator
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果接受 generator，則將 generator 新增到額外引數字典中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回準備好的額外引數字典
return extra_step_kwargs
# 檢查輸入引數的有效性
def check_inputs(
self,
prompt,
image,
original_image,
mask_image,
batch_size,
callback_steps,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if_img2img.IFImg2ImgPipeline.preprocess_image 複製而來，將 preprocess_image 修改為 preprocess_original_image
def preprocess_original_image(self, image: PIL.Image.Image) -> torch.Tensor:
# 如果輸入的 image 不是列表，則將其轉換為列表
if not isinstance(image, list):
image = [image]
# 定義將 NumPy 陣列轉換為 PyTorch 張量的函式
def numpy_to_pt(images):
# 如果輸入影像是 3D 陣列，則在最後新增一個維度
if images.ndim == 3:
images = images[..., None]
# 將 NumPy 陣列轉換為 PyTorch 張量並調整維度順序
images = torch.from_numpy(images.transpose(0, 3, 1, 2))
return images
# 如果第一個影像是 PIL 影像型別
if isinstance(image[0], PIL.Image.Image):
new_image = []
# 遍歷每個影像，進行轉換和處理
for image_ in image:
# 將影像轉換為 RGB 格式
image_ = image_.convert("RGB")
# 調整影像大小
image_ = resize(image_, self.unet.config.sample_size)
# 將影像轉換為 NumPy 陣列
image_ = np.array(image_)
# 轉換資料型別為 float32
image_ = image_.astype(np.float32)
# 歸一化影像資料到 [-1, 1] 範圍
image_ = image_ / 127.5 - 1
# 將處理後的影像新增到新影像列表中
new_image.append(image_)
# 將新影像列表轉換為 NumPy 陣列
image = new_image
# 將影像堆疊成一個 NumPy 陣列
image = np.stack(image, axis=0) # 轉換為 NumPy 陣列
# 將 NumPy 陣列轉換為 PyTorch 張量
image = numpy_to_pt(image) # 轉換為 PyTorch 張量
# 如果第一個影像是 NumPy 陣列
elif isinstance(image[0], np.ndarray):
# 如果陣列是 4 維，則進行拼接，否則堆疊成一個陣列
image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
# 將 NumPy 陣列轉換為 PyTorch 張量
image = numpy_to_pt(image)
# 如果第一個影像是 PyTorch 張量
elif isinstance(image[0], torch.Tensor):
# 如果張量是 4 維，則進行拼接，否則堆疊成一個張量
image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
# 返回處理後的影像張量
return image
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if_superresolution.IFSuperResolutionPipeline.preprocess_image 複製而來
# 定義影像預處理函式，接受影像、每個提示的影像數量和裝置作為引數，返回處理後的張量
def preprocess_image(self, image: PIL.Image.Image, num_images_per_prompt, device) -> torch.Tensor:
# 檢查輸入的影像是否為張量或列表，如果不是則將其轉換為列表
if not isinstance(image, torch.Tensor) and not isinstance(image, list):
image = [image]
# 如果列表中的第一個元素是 PIL 影像
if isinstance(image[0], PIL.Image.Image):
# 將 PIL 影像轉換為 NumPy 陣列，並歸一化到 [-1, 1] 範圍
image = [np.array(i).astype(np.float32) / 127.5 - 1.0 for i in image]
# 將列表中的影像堆疊為 NumPy 陣列，增加一個維度
image = np.stack(image, axis=0) # to np
# 將 NumPy 陣列轉換為 PyTorch 張量，並調整維度順序
image = torch.from_numpy(image.transpose(0, 3, 1, 2))
# 如果列表中的第一個元素是 NumPy 陣列
elif isinstance(image[0], np.ndarray):
# 將列表中的影像堆疊為 NumPy 陣列，增加一個維度
image = np.stack(image, axis=0) # to np
# 如果影像是 5 維，取第一個元素
if image.ndim == 5:
image = image[0]
# 將 NumPy 陣列轉換為 PyTorch 張量，並調整維度順序
image = torch.from_numpy(image.transpose(0, 3, 1, 2))
# 如果輸入是列表且第一個元素是張量
elif isinstance(image, list) and isinstance(image[0], torch.Tensor):
# 獲取第一個張量的維度
dims = image[0].ndim
# 如果是 3 維，沿第 0 維堆疊張量
if dims == 3:
image = torch.stack(image, dim=0)
# 如果是 4 維，沿第 0 維連線張量
elif dims == 4:
image = torch.concat(image, dim=0)
# 如果維度不是 3 或 4，丟擲錯誤
else:
raise ValueError(f"Image must have 3 or 4 dimensions, instead got {dims}")
# 將影像張量移動到指定裝置，並設定資料型別
image = image.to(device=device, dtype=self.unet.dtype)
# 重複影像以匹配每個提示的影像數量
image = image.repeat_interleave(num_images_per_prompt, dim=0)
# 返回處理後的影像張量
return image
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if_inpainting.IFInpaintingPipeline 複製的預處理掩碼影像的程式碼
# 預處理掩碼影像，返回處理後的 PyTorch 張量
def preprocess_mask_image(self, mask_image) -> torch.Tensor:
# 檢查掩碼影像是否為列表，如果不是，則將其包裝為列表
if not isinstance(mask_image, list):
mask_image = [mask_image]
# 如果掩碼影像的第一個元素是 PyTorch 張量
if isinstance(mask_image[0], torch.Tensor):
# 根據第一個張量的維度，選擇合併（cat）或堆疊（stack）操作
mask_image = torch.cat(mask_image, axis=0) if mask_image[0].ndim == 4 else torch.stack(mask_image, axis=0)
# 如果處理後的張量是二維
if mask_image.ndim == 2:
# 對於單個掩碼，增加批次維度和通道維度
mask_image = mask_image.unsqueeze(0).unsqueeze(0)
# 如果處理後的張量是三維，且第一維大小為1
elif mask_image.ndim == 3 and mask_image.shape[0] == 1:
# 單個掩碼，認為第0維是批次大小為1
mask_image = mask_image.unsqueeze(0)
# 如果處理後的張量是三維，且第一維大小不為1
elif mask_image.ndim == 3 and mask_image.shape[0] != 1:
# 一批掩碼，認為第0維是批次維度
mask_image = mask_image.unsqueeze(1)
# 將掩碼影像中小於0.5的值設為0
mask_image[mask_image < 0.5] = 0
# 將掩碼影像中大於等於0.5的值設為1
mask_image[mask_image >= 0.5] = 1
# 如果掩碼影像的第一個元素是 PIL 影像
elif isinstance(mask_image[0], PIL.Image.Image):
new_mask_image = [] # 建立一個新的列表以儲存處理後的掩碼影像
# 遍歷每個掩碼影像
for mask_image_ in mask_image:
# 將掩碼影像轉換為灰度模式
mask_image_ = mask_image_.convert("L")
# 調整掩碼影像的大小
mask_image_ = resize(mask_image_, self.unet.config.sample_size)
# 將影像轉換為 NumPy 陣列
mask_image_ = np.array(mask_image_)
# 增加批次和通道維度
mask_image_ = mask_image_[None, None, :]
# 將處理後的掩碼影像新增到新列表中
new_mask_image.append(mask_image_)
# 將新處理的掩碼影像列表合併為一個張量
mask_image = new_mask_image
# 沿第0維連線所有掩碼影像
mask_image = np.concatenate(mask_image, axis=0)
# 將值轉換為浮點數並歸一化到[0, 1]
mask_image = mask_image.astype(np.float32) / 255.0
# 將掩碼影像中小於0.5的值設為0
mask_image[mask_image < 0.5] = 0
# 將掩碼影像中大於等於0.5的值設為1
mask_image[mask_image >= 0.5] = 1
# 將 NumPy 陣列轉換為 PyTorch 張量
mask_image = torch.from_numpy(mask_image)
# 如果掩碼影像的第一個元素是 NumPy 陣列
elif isinstance(mask_image[0], np.ndarray):
# 沿第0維連線每個掩碼陣列，並增加批次和通道維度
mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0)
# 將掩碼影像中小於0.5的值設為0
mask_image[mask_image < 0.5] = 0
# 將掩碼影像中大於等於0.5的值設為1
mask_image[mask_image >= 0.5] = 1
# 將 NumPy 陣列轉換為 PyTorch 張量
mask_image = torch.from_numpy(mask_image)
# 返回處理後的掩碼影像
return mask_image
# 從 diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps 複製的程式碼
def get_timesteps(self, num_inference_steps, strength):
# 根據初始時間步計算原始時間步
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
# 計算開始時間步，確保不小於0
t_start = max(num_inference_steps - init_timestep, 0)
# 從排程器中獲取時間步
timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
# 如果排程器有設定開始索引的方法，則呼叫它
if hasattr(self.scheduler, "set_begin_index"):
self.scheduler.set_begin_index(t_start * self.scheduler.order)
# 返回時間步和剩餘的推理步驟數量
return timesteps, num_inference_steps - t_start
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if_inpainting.IFInpaintingPipeline.prepare_intermediate_images 複製的程式碼
# 準備中間影像，為影像處理生成噪聲並應用遮罩
def prepare_intermediate_images(
self, image, timestep, batch_size, num_images_per_prompt, dtype, device, mask_image, generator=None
):
# 獲取輸入影像的批次大小、通道數、高度和寬度
image_batch_size, channels, height, width = image.shape
# 更新批次大小為每個提示生成的影像數量
batch_size = batch_size * num_images_per_prompt
# 定義新的形狀用於生成噪聲
shape = (batch_size, channels, height, width)
# 檢查生成器列表的長度是否與請求的批次大小匹配
if isinstance(generator, list) and len(generator) != batch_size:
# 丟擲錯誤以確保生成器數量與批次大小一致
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
# 生成與輸入形狀匹配的隨機噪聲
noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 將輸入影像按提示的數量重複，以匹配新批次大小
image = image.repeat_interleave(num_images_per_prompt, dim=0)
# 向影像新增噪聲
noised_image = self.scheduler.add_noise(image, noise, timestep)
# 應用遮罩以混合原始影像和帶噪聲影像
image = (1 - mask_image) * image + mask_image * noised_image
# 返回處理後的影像
return image
# 禁用梯度計算以節省記憶體
@torch.no_grad()
# 替換示例文件字串
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
# 接受影像，支援多種輸入型別
image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
# 可選引數：原始影像，支援多種輸入型別
original_image: Union[
PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
] = None,
# 可選引數：遮罩影像，支援多種輸入型別
mask_image: Union[
PIL.Image.Image, torch.Tensor, np.ndarray, List[PIL.Image.Image], List[torch.Tensor], List[np.ndarray]
] = None,
# 設定強度引數，預設值為0.8
strength: float = 0.8,
# 可選引數：提示資訊，支援字串或字串列表
prompt: Union[str, List[str]] = None,
# 設定推理步驟的數量，預設值為100
num_inference_steps: int = 100,
# 可選引數：時間步列表
timesteps: List[int] = None,
# 設定指導尺度，預設值為4.0
guidance_scale: float = 4.0,
# 可選引數：負面提示資訊
negative_prompt: Optional[Union[str, List[str]]] = None,
# 可選引數：每個提示生成的影像數量，預設值為1
num_images_per_prompt: Optional[int] = 1,
# 設定η值，預設值為0.0
eta: float = 0.0,
# 可選引數：隨機數生成器
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 可選引數：提示嵌入
prompt_embeds: Optional[torch.Tensor] = None,
# 可選引數：負面提示嵌入
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 可選引數：輸出型別，預設值為"pil"
output_type: Optional[str] = "pil",
# 可選引數：是否返回字典，預設值為True
return_dict: bool = True,
# 可選引數：回撥函式
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 設定回撥步驟的數量，預設值為1
callback_steps: int = 1,
# 可選引數：交叉注意力關鍵字引數
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 設定噪聲水平，預設值為0
noise_level: int = 0,
# 可選引數：清理標題的標誌，預設值為True
clean_caption: bool = True,
.\diffusers\pipelines\deepfloyd_if\pipeline_if_superresolution.py
# 匯入用於處理 HTML 的模組
import html
# 匯入用於獲取物件資訊的模組
import inspect
# 匯入用於正規表示式處理的模組
import re
# 匯入 URL 解析模組並命名為 ul
import urllib.parse as ul
# 從 typing 模組匯入型別提示相關的類
from typing import Any, Callable, Dict, List, Optional, Union
# 匯入 NumPy 庫並命名為 np
import numpy as np
# 匯入影像處理庫 PIL
import PIL.Image
# 匯入 PyTorch 庫
import torch
# 匯入 PyTorch 的功能性模組
import torch.nn.functional as F
# 從 transformers 庫匯入影像處理器和 T5 模型及其標記器
from transformers import CLIPImageProcessor, T5EncoderModel, T5Tokenizer
# 從相對路徑匯入所需的混合類和模型
from ...loaders import StableDiffusionLoraLoaderMixin
from ...models import UNet2DConditionModel
from ...schedulers import DDPMScheduler
# 從 utils 模組匯入多個工具函式和常量
from ...utils import (
BACKENDS_MAPPING,
is_bs4_available,
is_ftfy_available,
logging,
replace_example_docstring,
)
# 從 torch_utils 模組匯入生成隨機張量的函式
from ...utils.torch_utils import randn_tensor
# 從 pipeline_utils 匯入擴散管道類
from ..pipeline_utils import DiffusionPipeline
# 從 pipeline_output 匯入輸出類
from .pipeline_output import IFPipelineOutput
# 從安全檢查器模組匯入安全檢查器類
from .safety_checker import IFSafetyChecker
# 從水印模組匯入水印處理類
from .watermark import IFWatermarker
# 如果 bs4 可用，則匯入 BeautifulSoup 類
if is_bs4_available():
from bs4 import BeautifulSoup
# 如果 ftfy 可用，則匯入該模組
if is_ftfy_available():
import ftfy
# 建立一個日誌記錄器以記錄模組內的日誌資訊
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# 示例文件字串，提供了使用該類的示例
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> from diffusers import IFPipeline, IFSuperResolutionPipeline, DiffusionPipeline
>>> from diffusers.utils import pt_to_pil
>>> import torch
>>> pipe = IFPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16)
>>> pipe.enable_model_cpu_offload()
>>> prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"'
>>> prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)
>>> image = pipe(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, output_type="pt").images
>>> # save intermediate image
>>> pil_image = pt_to_pil(image)
>>> pil_image[0].save("./if_stage_I.png")
>>> super_res_1_pipe = IFSuperResolutionPipeline.from_pretrained(
... "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16
... )
>>> super_res_1_pipe.enable_model_cpu_offload()
>>> image = super_res_1_pipe(
... image=image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds
... ).images
>>> image[0].save("./if_stage_II.png")
```py
"""
# 定義 IFSuperResolutionPipeline 類，繼承自 DiffusionPipeline 和 StableDiffusionLoraLoaderMixin
class IFSuperResolutionPipeline(DiffusionPipeline, StableDiffusionLoraLoaderMixin):
# 定義標記器屬性，型別為 T5Tokenizer
tokenizer: T5Tokenizer
# 定義文字編碼器屬性，型別為 T5EncoderModel
text_encoder: T5EncoderModel
# 定義 UNet 模型屬性，型別為 UNet2DConditionModel
unet: UNet2DConditionModel
# 定義排程器屬性，型別為 DDPMScheduler
scheduler: DDPMScheduler
# 定義影像噪聲排程器屬性，型別為 DDPMScheduler
image_noising_scheduler: DDPMScheduler
# 定義可選的特徵提取器屬性，型別為 CLIPImageProcessor
feature_extractor: Optional[CLIPImageProcessor]
# 定義可選的安全檢查器屬性，型別為 IFSafetyChecker
safety_checker: Optional[IFSafetyChecker]
# 定義可選的水印處理器屬性，型別為 IFWatermarker
watermarker: Optional[IFWatermarker]
# 定義用於匹配不良標點符號的正規表示式
bad_punct_regex = re.compile(
r"["
+ "#®•©™&@·º½¾¿¡§~"
+ r"\)"
+ r"\("
+ r"\]"
+ r"\["
+ r"\}"
+ r"\{"
+ r"\|"
+ "\\"
+ r"\/"
+ r"\*"
+ r"]{1,}"
) # noqa
# 定義可選元件的列表
_optional_components = ["tokenizer", "text_encoder", "safety_checker", "feature_extractor", "watermarker"]
# 定義模型在 CPU 上的解除安裝順序
model_cpu_offload_seq = "text_encoder->unet"
# 定義不參與 CPU 解除安裝的元件
_exclude_from_cpu_offload = ["watermarker"]
# 初始化方法，定義模型的各個元件
def __init__(
self,
tokenizer: T5Tokenizer, # 用於文字分詞的模型
text_encoder: T5EncoderModel, # 文字編碼器模型
unet: UNet2DConditionModel, # UNet 模型，用於影像生成
scheduler: DDPMScheduler, # 排程器，用於控制生成過程
image_noising_scheduler: DDPMScheduler, # 影像噪聲排程器
safety_checker: Optional[IFSafetyChecker], # 安全檢查器（可選）
feature_extractor: Optional[CLIPImageProcessor], # 特徵提取器（可選）
watermarker: Optional[IFWatermarker], # 水印器（可選）
requires_safety_checker: bool = True, # 是否需要安全檢查器的標誌
):
# 呼叫父類的初始化方法
super().__init__()
# 檢查安全檢查器是否為 None，且要求使用安全檢查器
if safety_checker is None and requires_safety_checker:
# 記錄警告資訊，提醒使用者安全檢查器未啟用
logger.warning(
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
" that you abide to the conditions of the IF license and do not expose unfiltered"
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
)
# 檢查安全檢查器存在且特徵提取器為 None
if safety_checker is not None and feature_extractor is None:
# 丟擲值錯誤，提醒使用者需要定義特徵提取器
raise ValueError(
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
)
# 檢查 UNet 的輸入通道數是否為 6
if unet.config.in_channels != 6:
# 記錄警告，提醒使用者所載入的檢查點不適合超解析度
logger.warning(
"It seems like you have loaded a checkpoint that shall not be used for super resolution from {unet.config._name_or_path} as it accepts {unet.config.in_channels} input channels instead of 6. Please make sure to pass a super resolution checkpoint as the `'unet'`: IFSuperResolutionPipeline.from_pretrained(unet=super_resolution_unet, ...)."
)
# 註冊各個元件到模型
self.register_modules(
tokenizer=tokenizer,
text_encoder=text_encoder,
unet=unet,
scheduler=scheduler,
image_noising_scheduler=image_noising_scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
watermarker=watermarker,
)
# 將是否需要安全檢查器的資訊註冊到配置中
self.register_to_config(requires_safety_checker=requires_safety_checker)
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing 複製的內容
# 定義文字預處理函式，接收文字和可選引數 clean_caption
def _text_preprocessing(self, text, clean_caption=False):
# 如果 clean_caption 為真且 bs4 庫不可用，則記錄警告資訊
if clean_caption and not is_bs4_available():
logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
# 記錄將 clean_caption 設定為 False 的警告資訊
logger.warning("Setting `clean_caption` to False...")
# 將 clean_caption 設定為 False
clean_caption = False
# 如果 clean_caption 為真且 ftfy 庫不可用，則記錄警告資訊
if clean_caption and not is_ftfy_available():
logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
# 記錄將 clean_caption 設定為 False 的警告資訊
logger.warning("Setting `clean_caption` to False...")
# 將 clean_caption 設定為 False
clean_caption = False
# 如果文字不是元組或列表，則將其轉換為列表
if not isinstance(text, (tuple, list)):
text = [text]
# 定義處理單個文字的內部函式
def process(text: str):
# 如果 clean_caption 為真，執行清理操作
if clean_caption:
text = self._clean_caption(text)
# 再次清理文字
text = self._clean_caption(text)
else:
# 將文字轉為小寫並去除首尾空白
text = text.lower().strip()
# 返回處理後的文字
return text
# 對文字列表中的每個文字應用處理函式並返回結果列表
return [process(t) for t in text]
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption 複製而來
@torch.no_grad()
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.encode_prompt 複製而來
def encode_prompt(
# 定義提示的型別為字串或字串列表
prompt: Union[str, List[str]],
# 是否使用無分類器自由引導，預設為 True
do_classifier_free_guidance: bool = True,
# 每個提示生成的影像數量，預設為 1
num_images_per_prompt: int = 1,
# 可選引數，指定裝置
device: Optional[torch.device] = None,
# 可選引數，負面提示，可以是字串或字串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 可選引數，提示嵌入
prompt_embeds: Optional[torch.Tensor] = None,
# 可選引數，負面提示嵌入
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 是否清理提示的選項，預設為 False
clean_caption: bool = False,
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.run_safety_checker 複製而來
def run_safety_checker(self, image, device, dtype):
# 如果存在安全檢查器
if self.safety_checker is not None:
# 使用特徵提取器處理影像並轉換為張量
safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)
# 執行安全檢查，返回處理後的影像和檢測結果
image, nsfw_detected, watermark_detected = self.safety_checker(
images=image,
clip_input=safety_checker_input.pixel_values.to(dtype=dtype),
)
else:
# 如果沒有安全檢查器，設定檢測結果為 None
nsfw_detected = None
watermark_detected = None
# 返回處理後的影像及其檢測結果
return image, nsfw_detected, watermark_detected
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_extra_step_kwargs 複製而來
# 準備額外的引數用於排程器步驟，因為並非所有排程器具有相同的引數簽名
def prepare_extra_step_kwargs(self, generator, eta):
# eta (η) 僅在 DDIMScheduler 中使用，其他排程器將忽略
# eta 對應於 DDIM 論文中的 η: https://arxiv.org/abs/2010.02502
# 應該在 [0, 1] 範圍內
# 檢查排程器的步驟是否接受 eta 引數
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 建立額外引數字典
extra_step_kwargs = {}
# 如果接受 eta，新增到額外引數中
if accepts_eta:
extra_step_kwargs["eta"] = eta
# 檢查排程器的步驟是否接受 generator 引數
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
# 如果接受 generator，新增到額外引數中
if accepts_generator:
extra_step_kwargs["generator"] = generator
# 返回包含額外引數的字典
return extra_step_kwargs
# 檢查輸入引數的有效性
def check_inputs(
self,
prompt,
image,
batch_size,
noise_level,
callback_steps,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
# 從 diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline.prepare_intermediate_images 複製
def prepare_intermediate_images(self, batch_size, num_channels, height, width, dtype, device, generator):
# 定義中間影像的形狀
shape = (batch_size, num_channels, height, width)
# 檢查生成器列表的長度是否與請求的批大小匹配
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"您傳入了長度為 {len(generator)} 的生成器列表，但請求的有效批大小為 {batch_size}。"
f" 請確保批大小與生成器的長度匹配。"
)
# 生成隨機噪聲張量作為初始中間影像
intermediate_images = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
# 按排程器所需的標準差縮放初始噪聲
intermediate_images = intermediate_images * self.scheduler.init_noise_sigma
# 返回生成的中間影像
return intermediate_images
# 預處理影像，使其適合於後續模型處理
def preprocess_image(self, image, num_images_per_prompt, device):
# 檢查輸入是否為張量或列表，若不是，則將其轉換為列表
if not isinstance(image, torch.Tensor) and not isinstance(image, list):
image = [image]
# 如果輸入影像是 PIL 影像，則轉換為 NumPy 陣列並歸一化
if isinstance(image[0], PIL.Image.Image):
image = [np.array(i).astype(np.float32) / 127.5 - 1.0 for i in image]
# 將影像列表堆疊為 NumPy 陣列
image = np.stack(image, axis=0) # to np
# 轉換 NumPy 陣列為 PyTorch 張量，並調整維度順序
image = torch.from_numpy(image.transpose(0, 3, 1, 2))
# 如果輸入影像是 NumPy 陣列，則直接堆疊並處理
elif isinstance(image[0], np.ndarray):
image = np.stack(image, axis=0) # to np
# 如果陣列有五個維度，則只取第一個
if image.ndim == 5:
image = image[0]
# 轉換 NumPy 陣列為 PyTorch 張量，並調整維度順序
image = torch.from_numpy(image.transpose(0, 3, 1, 2))
# 如果輸入是張量列表，則根據維度進行堆疊或連線
elif isinstance(image, list) and isinstance(image[0], torch.Tensor):
dims = image[0].ndim
# 三維張量則堆疊，四維張量則連線
if dims == 3:
image = torch.stack(image, dim=0)
elif dims == 4:
image = torch.concat(image, dim=0)
# 若維度不符合要求，丟擲錯誤
else:
raise ValueError(f"Image must have 3 or 4 dimensions, instead got {dims}")
# 將影像移動到指定裝置，並設定資料型別
image = image.to(device=device, dtype=self.unet.dtype)
# 根據每個提示生成的影像數量重複影像
image = image.repeat_interleave(num_images_per_prompt, dim=0)
# 返回處理後的影像
return image
# 裝飾器，用於禁用梯度計算，節省記憶體和計算
@torch.no_grad()
# 替換示例文件字串的裝飾器
@replace_example_docstring(EXAMPLE_DOC_STRING)
# 定義可呼叫的類方法
def __call__(
# 提示內容，可以是字串或字串列表
prompt: Union[str, List[str]] = None,
# 輸出影像的高度
height: int = None,
# 輸出影像的寬度
width: int = None,
# 輸入的影像，可以是多種型別
image: Union[PIL.Image.Image, np.ndarray, torch.Tensor] = None,
# 推理步驟的數量
num_inference_steps: int = 50,
# 時間步列表
timesteps: List[int] = None,
# 指導縮放因子
guidance_scale: float = 4.0,
# 負提示，可以是字串或字串列表
negative_prompt: Optional[Union[str, List[str]]] = None,
# 每個提示生成的影像數量
num_images_per_prompt: Optional[int] = 1,
# 噪聲級別
eta: float = 0.0,
# 隨機數生成器，可以是單個或列表
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
# 提示嵌入張量
prompt_embeds: Optional[torch.Tensor] = None,
# 負提示嵌入張量
negative_prompt_embeds: Optional[torch.Tensor] = None,
# 輸出型別，預設為 PIL
output_type: Optional[str] = "pil",
# 是否返回字典格式
return_dict: bool = True,
# 可選的回撥函式
callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
# 回撥步驟的數量
callback_steps: int = 1,
# 交叉注意力的引數
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
# 噪聲級別
noise_level: int = 250,
# 是否清理標題
clean_caption: bool = True,
.\diffusers\pipelines\deepfloyd_if\pipeline_output.py
# 從 dataclasses 模組匯入 dataclass 裝飾器，用於簡化類的定義
from dataclasses import dataclass
# 從 typing 模組匯入型別註解，用於型別提示
from typing import List, Optional, Union
# 匯入 numpy 庫，通常用於陣列操作和數值計算
import numpy as np
# 匯入 PIL.Image，用於處理影像
import PIL.Image
# 從上級模組匯入 BaseOutput 基類，用於輸出類的繼承
from ...utils import BaseOutput
# 定義 IFPipelineOutput 類，繼承自 BaseOutput
@dataclass
class IFPipelineOutput(BaseOutput):
"""
Args:
Output class for Stable Diffusion pipelines.
images (`List[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
nsfw_detected (`List[bool]`)
List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
(nsfw) content or a watermark. `None` if safety checking could not be performed.
watermark_detected (`List[bool]`)
List of flags denoting whether the corresponding generated image likely has a watermark. `None` if safety
checking could not be performed.
"""
# 定義 images 屬性，可以是 PIL 影像列表或 numpy 陣列
images: Union[List[PIL.Image.Image], np.ndarray]
# 定義 nsfw_detected 屬性，可選的布林列表，用於標記是否檢測到不安全內容
nsfw_detected: Optional[List[bool]]
# 定義 watermark_detected 屬性，可選的布林列表，用於標記是否檢測到水印
watermark_detected: Optional[List[bool]]
.\diffusers\pipelines\deepfloyd_if\safety_checker.py
# 匯入必要的庫和模組
import numpy as np # 匯入 NumPy 庫，用於數值計算
import torch # 匯入 PyTorch 庫，用於深度學習
import torch.nn as nn # 匯入 PyTorch 的神經網路模組
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel # 匯入 CLIP 配置和模型
from ...utils import logging # 從父級模組匯入日誌工具
# 獲取當前模組的日誌記錄器
logger = logging.get_logger(__name__)
# 定義一個安全檢查器類，繼承自預訓練模型
class IFSafetyChecker(PreTrainedModel):
# 指定配置類為 CLIPConfig
config_class = CLIPConfig
# 指定不進行拆分的模組
_no_split_modules = ["CLIPEncoderLayer"]
# 初始化方法，接受一個配置物件
def __init__(self, config: CLIPConfig):
# 呼叫父類的初始化方法
super().__init__(config)
# 建立視覺模型，使用配置中的視覺部分
self.vision_model = CLIPVisionModelWithProjection(config.vision_config)
# 定義線性層用於 NSFW 檢測
self.p_head = nn.Linear(config.vision_config.projection_dim, 1)
# 定義線性層用於水印檢測
self.w_head = nn.Linear(config.vision_config.projection_dim, 1)
# 無梯度計算的前向傳播方法
@torch.no_grad()
def forward(self, clip_input, images, p_threshold=0.5, w_threshold=0.5):
# 獲取影像的嵌入向量
image_embeds = self.vision_model(clip_input)[0]
# 檢測 NSFW 內容
nsfw_detected = self.p_head(image_embeds)
# 將輸出展平為一維
nsfw_detected = nsfw_detected.flatten()
# 根據閾值判斷是否檢測到 NSFW 內容
nsfw_detected = nsfw_detected > p_threshold
# 轉換為列表格式
nsfw_detected = nsfw_detected.tolist()
# 如果檢測到 NSFW 內容，記錄警告日誌
if any(nsfw_detected):
logger.warning(
"Potential NSFW content was detected in one or more images. A black image will be returned instead."
" Try again with a different prompt and/or seed."
)
# 遍歷每個影像，處理檢測到 NSFW 內容的影像
for idx, nsfw_detected_ in enumerate(nsfw_detected):
if nsfw_detected_:
# 將檢測到的 NSFW 影像替換為全黑影像
images[idx] = np.zeros(images[idx].shape)
# 檢測水印內容
watermark_detected = self.w_head(image_embeds)
# 將輸出展平為一維
watermark_detected = watermark_detected.flatten()
# 根據閾值判斷是否檢測到水印內容
watermark_detected = watermark_detected > w_threshold
# 轉換為列表格式
watermark_detected = watermark_detected.tolist()
# 如果檢測到水印內容，記錄警告日誌
if any(watermark_detected):
logger.warning(
"Potential watermarked content was detected in one or more images. A black image will be returned instead."
" Try again with a different prompt and/or seed."
)
# 遍歷每個影像，處理檢測到水印的影像
for idx, watermark_detected_ in enumerate(watermark_detected):
if watermark_detected_:
# 將檢測到水印的影像替換為全黑影像
images[idx] = np.zeros(images[idx].shape)
# 返回處理後的影像和檢測結果
return images, nsfw_detected, watermark_detected
.\diffusers\pipelines\deepfloyd_if\timesteps.py
# 定義一個包含 27 個時間步長的列表，表示快速模式下的時間點
fast27_timesteps = [
# 具體的時間步長值
999,
800,
799,
600,
599,
500,
400,
399,
377,
355,
333,
311,
288,
266,
244,
222,
200,
199,
177,
155,
133,
111,
88,
66,
44,
22,
0,
]
# 定義一個包含 27 個時間步長的列表，表示智慧模式下的時間點
smart27_timesteps = [
# 具體的時間步長值
999,
976,
952,
928,
905,
882,
858,
857,
810,
762,
715,
714,
572,
429,
428,
286,
285,
238,
190,
143,
142,
118,
95,
71,
47,
24,
0,
]
# 定義一個包含 100 個時間步長的列表，表示智慧模式下的時間點
smart50_timesteps = [
# 具體的時間步長值
999,
988,
977,
966,
955,
944,
933,
922,
911,
900,
899,
879,
859,
840,
820,
800,
799,
766,
733,
700,
699,
650,
600,
599,
500,
499,
400,
399,
350,
300,
299,
266,
233,
200,
199,
179,
159,
140,
120,
100,
99,
88,
77,
66,
55,
44,
33,
22,
11,
0,
]
# 定義一個包含 185 個時間步長的列表，表示智慧模式下的時間點
smart100_timesteps = [
# 具體的時間步長值
999,
995,
992,
989,
985,
981,
978,
975,
971,
967,
964,
961,
957,
956,
951,
947,
942,
937,
933,
928,
923,
919,
914,
913,
908,
903,
897,
892,
887,
881,
876,
871,
870,
864,
858,
852,
846,
840,
834,
828,
827,
820,
813,
806,
799,
792,
785,
784,
777,
770,
763,
756,
749,
742,
741,
733,
724,
716,
707,
699,
698,
688,
677,
666,
656,
655,
645,
634,
623,
613,
612,
598,
584,
570,
569,
555,
541,
527,
526,
505,
484,
483,
462,
440,
439,
396,
395,
352,
351,
308,
307,
264,
263,
220,
219,
176,
132,
88,
44,
0,
]
# 定義一個包含 185 個時間步長的列表，表示智慧模式下的時間點
smart185_timesteps = [
# 具體的時間步長值
999,
997,
995,
992,
990,
988,
986,
984,
981,
979,
977,
975,
972,
970,
968,
966,
964,
961,
959,
957,
956,
954,
951,
949,
946,
944,
941,
939,
936,
934,
931,
929,
926,
924,
921,
919,
916,
914,
913,
910,
907,
905,
902,
899,
896,
893,
891,
888,
885,
882,
879,
877,
874,
871,
870,
867,
864,
861,
858,
855,
852,
849,
846,
843,
840,
837,
834,
831,
828,
827,
824,
821,
817,
814,
811,
808,
804,
801,
798,
795,
791,
788,
785,
784,
780,
777,
774,
770,
766,
763,
760,
756,
752,
749,
746,
742,
741,
737,
733,
730,
726,
722,
718,
714,
710,
707,
703,
699,
698,
694,
690,
685,
681,
677,
673,
669,
664,
660,
# 新增數字 656 到列表
656,
# 新增數字 655 到列表
655,
# 新增數字 650 到列表
650,
# 新增數字 646 到列表
646,
# 新增數字 641 到列表
641,
# 新增數字 636 到列表
636,
# 新增數字 632 到列表
632,
# 新增數字 627 到列表
627,
# 新增數字 622 到列表
622,
# 新增數字 618 到列表
618,
# 新增數字 613 到列表
613,
# 新增數字 612 到列表
612,
# 新增數字 607 到列表
607,
# 新增數字 602 到列表
602,
# 新增數字 596 到列表
596,
# 新增數字 591 到列表
591,
# 新增數字 586 到列表
586,
# 新增數字 580 到列表
580,
# 新增數字 575 到列表
575,
# 新增數字 570 到列表
570,
# 新增數字 569 到列表
569,
# 新增數字 563 到列表
563,
# 新增數字 557 到列表
557,
# 新增數字 551 到列表
551,
# 新增數字 545 到列表
545,
# 新增數字 539 到列表
539,
# 新增數字 533 到列表
533,
# 新增數字 527 到列表
527,
# 新增數字 526 到列表
526,
# 新增數字 519 到列表
519,
# 新增數字 512 到列表
512,
# 新增數字 505 到列表
505,
# 新增數字 498 到列表
498,
# 新增數字 491 到列表
491,
# 新增數字 484 到列表
484,
# 新增數字 483 到列表
483,
# 新增數字 474 到列表
474,
# 新增數字 466 到列表
466,
# 新增數字 457 到列表
457,
# 新增數字 449 到列表
449,
# 新增數字 440 到列表
440,
# 新增數字 439 到列表
439,
# 新增數字 428 到列表
428,
# 新增數字 418 到列表
418,
# 新增數字 407 到列表
407,
# 新增數字 396 到列表
396,
# 新增數字 395 到列表
395,
# 新增數字 381 到列表
381,
# 新增數字 366 到列表
366,
# 新增數字 352 到列表
352,
# 新增數字 351 到列表
351,
# 新增數字 330 到列表
330,
# 新增數字 308 到列表
308,
# 新增數字 307 到列表
307,
# 新增數字 286 到列表
286,
# 新增數字 264 到列表
264,
# 新增數字 263 到列表
263,
# 新增數字 242 到列表
242,
# 新增數字 220 到列表
220,
# 新增數字 219 到列表
219,
# 新增數字 176 到列表
176,
# 新增數字 175 到列表
175,
# 新增數字 132 到列表
132,
# 新增數字 131 到列表
131,
# 新增數字 88 到列表
88,
# 新增數字 44 到列表
44,
# 新增數字 0 到列表
0,
# 定義一個空列表，後續將填充時間步長資料
super27_timesteps = [
# 定義時間步長，值遞減
999,
991,
982,
974,
966,
958,
950,
941,
933,
925,
916,
908,
900,
899,
874,
850,
825,
800,
799,
700,
600,
500,
400,
300,
200,
100,
0,
]
# 定義另一個時間步長列表，值同樣遞減
super40_timesteps = [
# 各時間步長的值
999,
992,
985,
978,
971,
964,
957,
949,
942,
935,
928,
921,
914,
907,
900,
899,
879,
859,
840,
820,
800,
799,
766,
733,
700,
699,
650,
600,
599,
500,
499,
400,
399,
300,
299,
200,
199,
100,
99,
0,
]
# 定義第三個時間步長列表，值繼續遞減
super100_timesteps = [
# 包含一系列時間步長的值
999,
996,
992,
989,
985,
982,
979,
975,
972,
968,
965,
961,
958,
955,
951,
948,
944,
941,
938,
934,
931,
927,
924,
920,
917,
914,
910,
907,
903,
900,
899,
891,
884,
876,
869,
861,
853,
846,
838,
830,
823,
815,
808,
800,
799,
788,
777,
766,
755,
744,
733,
722,
711,
700,
699,
688,
677,
666,
655,
644,
633,
622,
611,
600,
599,
585,
571,
557,
542,
528,
514,
500,
499,
485,
471,
457,
442,
428,
414,
400,
399,
379,
359,
340,
320,
300,
299,
279,
259,
240,
220,
200,
199,
166,
133,
100,
99,
66,
33,
0,
]
.\diffusers\pipelines\deepfloyd_if\watermark.py
# 從 typing 模組匯入 List 型別，用於型別註釋
from typing import List
# 匯入 PIL.Image 庫以處理影像
import PIL.Image
# 匯入 torch 庫用於張量操作
import torch
# 從 PIL 匯入 Image 類以建立和處理影像
from PIL import Image
# 從配置工具模組匯入 ConfigMixin 類
from ...configuration_utils import ConfigMixin
# 從模型工具模組匯入 ModelMixin 類
from ...models.modeling_utils import ModelMixin
# 從工具模組匯入 PIL_INTERPOLATION 以獲取插值方法
from ...utils import PIL_INTERPOLATION
# 定義 IFWatermarker 類，繼承自 ModelMixin 和 ConfigMixin
class IFWatermarker(ModelMixin, ConfigMixin):
# 初始化方法
def __init__(self):
# 呼叫父類初始化方法
super().__init__()
# 註冊一個形狀為 (62, 62, 4) 的零張量作為水印影像
self.register_buffer("watermark_image", torch.zeros((62, 62, 4)))
# 初始化水印影像的 PIL 表示為 None
self.watermark_image_as_pil = None
# 定義應用水印的方法，接受影像列表和可選的樣本大小
def apply_watermark(self, images: List[PIL.Image.Image], sample_size=None):
# 從 GitHub 複製的程式碼
# 獲取第一張影像的高度
h = images[0].height
# 獲取第一張影像的寬度
w = images[0].width
# 如果未指定樣本大小，則使用影像高度
sample_size = sample_size or h
# 計算寬高比係數
coef = min(h / sample_size, w / sample_size)
# 根據係數計算影像的新高度和寬度
img_h, img_w = (int(h / coef), int(w / coef)) if coef < 1 else (h, w)
# 定義 S1 和 S2，用於計算 K
S1, S2 = 1024**2, img_w * img_h
# 計算 K 值
K = (S2 / S1) ** 0.5
# 計算水印大小及其在影像中的位置
wm_size, wm_x, wm_y = int(K * 62), img_w - int(14 * K), img_h - int(14 * K)
# 如果水印影像尚未建立
if self.watermark_image_as_pil is None:
# 將水印張量轉換為 uint8 型別並轉移到 CPU，轉換為 NumPy 陣列
watermark_image = self.watermark_image.to(torch.uint8).cpu().numpy()
# 將 NumPy 陣列轉換為 RGBA 模式的 PIL 影像
watermark_image = Image.fromarray(watermark_image, mode="RGBA")
# 將 PIL 影像儲存到例項變數
self.watermark_image_as_pil = watermark_image
# 調整水印影像大小
wm_img = self.watermark_image_as_pil.resize(
(wm_size, wm_size), PIL_INTERPOLATION["bicubic"], reducing_gap=None
)
# 遍歷輸入影像列表
for pil_img in images:
# 將水印影像貼上到每張影像上，使用水印的 alpha 通道作為掩碼
pil_img.paste(wm_img, box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y), mask=wm_img.split()[-1])
# 返回新增水印後的影像列表
return images