FLUX.1 图生图原理及实践

MODEL ZOO Oct 28, 2024

这篇文章将指导你根据现有图像和文本提示生成新图像。这项技术在一篇名为《SDEdit：使用随机微分方程进行引导图像合成和编辑》的论文中提出，并在此应用于 FLUX.1。

首先，我们将简要解释潜在扩散模型的工作原理。然后，我们将了解 SDEdit 如何修改反向扩散过程以根据文本提示编辑图像。最后，我们将提供运行整个管道的代码。

1、背景：潜在扩散

潜在扩散在低维潜在空间中执行扩散过程。让我们定义潜在空间：

Source: https://en.wikipedia.org/wiki/Variational_autoencoder

变分自动编码器 (VAE) 将图像从像素空间（人类理解的 RGB 高度宽度表示）投影到较小的潜在空间。这种压缩保留了足够的信息以便稍后重建图像。扩散过程在此潜在空间中运行，因为它在计算上更便宜，并且对不相关的像素空间细节不太敏感。

现在，让我们解释一下潜在扩散：

Source: https://en.wikipedia.org/wiki/Diffusion_model

扩散过程分为两部分：

前向扩散：一种预定的、非学习的过程，通过多个步骤将自然图像转换为纯噪声。
后向扩散：一种学习过程，从纯噪声重建自然图像。

请注意，噪声被添加到潜在空间并遵循特定的时间表，在前向过程中从弱到强。

噪声按照特定的时间表添加到潜在空间，在前向扩散过程中从弱噪声逐渐变为强噪声。与 GAN 等一次性生成方法相比，这种多步骤方法简化了网络的任务。反向过程是通过似然最大化来学习的，这比对抗性损失更容易优化。

文本条件

Source: https://github.com/CompVis/latent-diffusion

图像生成还取决于文本等额外信息，这是你可能提供给稳定扩散或 Flux.1 模型的提示。在学习如何进行反向过程时，此文本作为扩散模型的“提示”包含在内。此文本使用类似 CLIP 或 T5 模型的模型进行编码，并输入到 UNet 或 Transformer 中，以将其引导至受噪声干扰的正确原始图像。

2、SDEdit

SDEdit 背后的想法很简单：在反向过程中，它不是像上图的“步骤 1”那样从完全随机噪声开始，而是从输入图像 + 缩放的随机噪声开始，然后再运行常规的反向扩散过程。因此，它如下：

加载输入图像，为 VAE 进行预处理
通过 VAE 运行它并采样一个输出（VAE 返回一个分布，因此我们需要采样来获取分布的一个实例）。
选择反向扩散过程的起始步骤 t_i。
采样一些缩放到 t_i 级别的噪声并将其添加到清晰图像表示中。
使用嘈杂的清晰图像和提示从 t_i 开始反向扩散过程。
使用 VAE 将结果投影回像素空间。
成功！

3、代码实现

以下是使用diffusers运行此工作流程的方法：

首先，安装依赖项：

pip install git+https://github.com/huggingface/diffusers.git optimum-quanto

目前，你需要从源代码安装diffusers，因为此功能在 pypi 上尚不可用。

接下来，加载 FluxImg2Img 管道：

import os

from diffusers import FluxImg2ImgPipeline
from optimum.quanto import qint8, qint4, quantize, freeze
import torch
from typing import Callable, List, Optional, Union, Dict, Any

from PIL import Image
import requests
import io

MODEL_PATH = os.getenv("MODEL_PATH", "black-forest-labs/FLUX.1-dev")

pipeline = FluxImg2ImgPipeline.from_pretrained(MODEL_PATH, torch_dtype=torch.bfloat16)

quantize(pipeline.text_encoder, weights=qint4, exclude="proj_out")
freeze(pipeline.text_encoder)

quantize(pipeline.text_encoder_2, weights=qint4, exclude="proj_out")
freeze(pipeline.text_encoder_2)

quantize(pipeline.transformer, weights=qint8, exclude="proj_out")
freeze(pipeline.transformer)

pipeline = pipeline.to("cuda")

generator = torch.Generator(device="cuda").manual_seed(100)

此代码加载管道并量化其中的某些部分，以便它适合 Colab 上可用的 L4 GPU。

现在，让我们定义一个实用函数来加载正确大小且不失真的图像：

def resize_image_center_crop(image_path_or_url, target_width, target_height):
    """
    Resizes an image while maintaining aspect ratio using center cropping.
    Handles both local file paths and URLs.

    Args:
        image_path_or_url: Path to the image file or URL.
        target_width: Desired width of the output image.
        target_height: Desired height of the output image.

    Returns:
        A PIL Image object with the resized image, or None if there's an error.
    """
    try:
        if image_path_or_url.startswith(('http://', 'https://')):  # Check if it's a URL
            response = requests.get(image_path_or_url, stream=True)
            response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
            img = Image.open(io.BytesIO(response.content))
        else:  # Assume it's a local file path
            img = Image.open(image_path_or_url)

        img_width, img_height = img.size

        # Calculate aspect ratios
        aspect_ratio_img = img_width / img_height
        aspect_ratio_target = target_width / target_height

        # Determine cropping box
        if aspect_ratio_img > aspect_ratio_target:  # Image is wider than target
            new_width = int(img_height * aspect_ratio_target)
            left = (img_width - new_width) // 2
            right = left + new_width
            top = 0
            bottom = img_height
        else:  # Image is taller or equal to target
            new_height = int(img_width / aspect_ratio_target)
            left = 0
            right = img_width
            top = (img_height - new_height) // 2
            bottom = top + new_height

        # Crop the image
        cropped_img = img.crop((left, top, right, bottom))

        # Resize to target dimensions
        resized_img = cropped_img.resize((target_width, target_height), Image.LANCZOS)

        return resized_img

    except (FileNotFoundError, requests.exceptions.RequestException, IOError) as e:
        print(f"Error: Could not open or process image from '{image_path_or_url}'.  Error: {e}")
        return None
    except Exception as e: #Catch other potential exceptions during image processing.
        print(f"An unexpected error occurred: {e}")
        return None

最后，让我们加载图像并运行管道：

url = "https://images.unsplash.com/photo-1609665558965-8e4c789cd7c5?ixlib=rb-4.0.3&q=85&fm=jpg&crop=entropy&cs=srgb&dl=sven-mieke-G-8B32scqMc-unsplash.jpg"
image = resize_image_center_crop(image_path_or_url=url, target_width=1024, target_height=1024)

prompt = "A picture of a Tiger"
image2 = pipeline(prompt, image=image, guidance_scale=3.5, generator=generator, height=1024, width=1024, num_inference_steps=28, strength=0.9).images[0]

这将下面的图像：