dreamfuse_inference.py

import gc
import os
from typing import List
import contextlib
import torch.multiprocessing as mp
from dataclasses import dataclass, field
from collections import defaultdict
import random
import numpy as np
from PIL import Image, ImageOps
import json
import torch
from peft import PeftModel
import torch.nn.functional as F
import accelerate
import diffusers
from diffusers import FluxPipeline
from diffusers.utils.torch_utils import is_compiled_module
import transformers
from tqdm import tqdm
from peft import LoraConfig, set_peft_model_state_dict
from peft.utils import get_peft_model_state_dict
from dreamfuse.models.dreamfuse_flux.transformer import (
    FluxTransformer2DModel,
    FluxTransformerBlock,
    FluxSingleTransformerBlock,
)
from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
    FlowMatchEulerDiscreteScheduler,
)
from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
from dreamfuse.trains.utils.inference_utils import (
    compute_text_embeddings,
    prepare_latents,
    _unpack_latents,
    _pack_latents,
    _prepare_image_ids,
    encode_images_cond,
    get_mask_affine,
    warp_affine_tensor
)


def seed_everything(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

@dataclass
class InferenceConfig:
    # Model paths
    flux_model_id: str = 'black-forest-labs/FLUX.1-dev'
    
    lora_id: str = ''
    model_choice: str = 'dev'
    # Model configs
    lora_rank: int = 16
    max_sequence_length: int = 256
    guidance_scale: float = 3.5
    num_inference_steps: int = 28
    mask_ids: int = 16
    mask_in_chans: int = 128 
    mask_out_chans: int = 3072
    inference_scale = 1024
    
    # Training configs
    gradient_checkpointing: bool = False
    mix_attention_double: bool = True
    mix_attention_single: bool = True
    
    # Image processing
    image_ids_offset: List[int] = field(default_factory=lambda: [0, 0, 0])
    image_tags: List[int] = field(default_factory=lambda: [0, 1, 2])
    context_tags: List[int] = None
    
    # Runtime configs
    device: str = "cuda:0" # if torch.cuda.is_available() else "cpu"
    dtype: torch.dtype = torch.bfloat16
    seed: int = 1234
    debug: bool = True

    # I/O configs
    valid_output_dir: str = "./inference_output"
    valid_roots: List[str] = field(default_factory=lambda: [
        "./",
    ])
    valid_jsons: List[str] = field(default_factory=lambda: [
        "./examples/data_dreamfuse.json",
    ])
    ref_prompts: str = ""
    
    truecfg: bool = False
    text_strength: int = 5
    
    # multi gpu
    sub_idx:int = 0
    total_num:int = 1

def adjust_fg_to_bg(image: Image.Image, mask: Image.Image, target_size: tuple) -> tuple[Image.Image, Image.Image]:
    width, height = image.size
    target_w, target_h = target_size
    
    scale = min(target_w / width, target_h / height)
    if scale < 1:
        new_w = int(width * scale)
        new_h = int(height * scale)
        image = image.resize((new_w, new_h))
        mask = mask.resize((new_w, new_h))
        width, height = new_w, new_h
    
    pad_w = target_w - width
    pad_h = target_h - height
    padding = (
        pad_w // 2,  # left
        pad_h // 2,  # top 
        (pad_w + 1) // 2,  # right
        (pad_h + 1) // 2   # bottom
    )
    
    image = ImageOps.expand(image, border=padding, fill=(255, 255, 255))
    mask = ImageOps.expand(mask, border=padding, fill=0)
    
    return image, mask

def find_nearest_bucket_size(input_width, input_height, mode="x64", bucket_size=1024):
    """
    Finds the nearest bucket size for the given input size.
    """
    buckets = {
        512: [[ 256, 768 ], [ 320, 768 ], [ 320, 704 ], [ 384, 640 ], [ 448, 576 ], [ 512, 512 ], [ 576, 448 ], [ 640, 384 ], [ 704, 320 ], [ 768, 320 ], [ 768, 256 ]],
        768: [[ 384, 1152 ], [ 480, 1152 ], [ 480, 1056 ], [ 576, 960 ], [ 672, 864 ], [ 768, 768 ], [ 864, 672 ], [ 960, 576 ], [ 1056, 480 ], [ 1152, 480 ], [ 1152, 384 ]],
        1024: [[ 512, 1536 ], [ 640, 1536 ], [ 640, 1408 ], [ 768, 1280 ], [ 896, 1152 ], [ 1024, 1024 ], [ 1152, 896 ], [ 1280, 768 ], [ 1408, 640 ], [ 1536, 640 ], [ 1536, 512 ]]
    }
    
    buckets = buckets[bucket_size]

    aspect_ratios = [w / h for (w, h) in buckets]
    assert mode in ["x64", "x8"]
    if mode == "x64":
        asp = input_width / input_height
        diff = [abs(ar - asp) for ar in aspect_ratios]
        bucket_id = int(np.argmin(diff))
        gen_width, gen_height = buckets[bucket_id]
    elif mode == "x8":
        max_pixels = 1024 * 1024
        ratio = (max_pixels / (input_width * input_height)) ** (0.5)
        gen_width, gen_height = round(input_width * ratio), round(input_height * ratio)
        gen_width = gen_width - gen_width % 8
        gen_height = gen_height - gen_height % 8
    else:
        raise NotImplementedError
    return (gen_width, gen_height)

def make_image_grid(images, rows, cols, size=None):
    assert len(images) == rows * cols

    if size is not None:
        images = [img.resize((size[0], size[1])) for img in images]

    w, h = images[0].size
    grid = Image.new("RGB", size=(cols * w, rows * h))

    for i, img in enumerate(images):
        grid.paste(img.convert("RGB"), box=(i % cols * w, i // cols * h))
    return grid

class DreamFuseInference:
    def __init__(self, config: InferenceConfig):
        self.config = config
        self.device = torch.device(config.device)
        torch.backends.cuda.matmul.allow_tf32 = True
        seed_everything(config.seed)
        self._init_models()

    def _init_models(self):
        # Initialize tokenizers
        self.tokenizer_one = transformers.CLIPTokenizer.from_pretrained(
            self.config.flux_model_id, subfolder="tokenizer"
        )
        self.tokenizer_two = transformers.T5TokenizerFast.from_pretrained(
            self.config.flux_model_id, subfolder="tokenizer_2"
        )

        # Initialize text encoders
        self.text_encoder_one = transformers.CLIPTextModel.from_pretrained(
            self.config.flux_model_id, subfolder="text_encoder"
        ).to(device=self.device, dtype=self.config.dtype)
        self.text_encoder_two = transformers.T5EncoderModel.from_pretrained(
            self.config.flux_model_id, subfolder="text_encoder_2"
        ).to(device=self.device, dtype=self.config.dtype)

        # Initialize VAE
        self.vae = diffusers.AutoencoderKL.from_pretrained(
            self.config.flux_model_id, subfolder="vae"
        ).to(device=self.device, dtype=self.config.dtype)

        # Initialize denoising model
        self.denoise_model = FluxTransformer2DModel.from_pretrained(
            self.config.flux_model_id, subfolder="transformer"
        ).to(device=self.device, dtype=self.config.dtype)

        if self.config.image_tags is not None or self.config.context_tags is not None:
            num_image_tag_embeddings = max(self.config.image_tags) + 1 if self.config.image_tags is not None else 0
            num_context_tag_embeddings = max(self.config.context_tags) + 1 if self.config.context_tags is not None else 0
            self.denoise_model.set_tag_embeddings(
                num_image_tag_embeddings=num_image_tag_embeddings,
                num_context_tag_embeddings=num_context_tag_embeddings,
            )

        # Add LoRA
        self.denoise_model = PeftModel.from_pretrained(
            self.denoise_model,
            self.config.lora_id,
            adapter_weights=[1.0],
            device_map={"": self.device}
        )

        # Initialize scheduler
        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
            self.config.flux_model_id, subfolder="scheduler"
        )

        # Set models to eval mode
        for model in [self.text_encoder_one, self.text_encoder_two, self.vae, self.denoise_model]:
            model.eval()
            model.requires_grad_(False)

    def _compute_text_embeddings(self, prompt):
        return compute_text_embeddings(
            self.config,
            prompt,
            [self.text_encoder_one, self.text_encoder_two],
            [self.tokenizer_one, self.tokenizer_two],
            self.device
        )

    def resize_to_fit_within(self, reference_image, target_image):
        ref_width, ref_height = reference_image.size
        target_width, target_height = target_image.size
        
        scale_width = ref_width / target_width
        scale_height = ref_height / target_height
        scale = min(scale_width, scale_height) 
        
        new_width = int(target_width * scale)
        new_height = int(target_height * scale)
        
        resized_image = target_image.resize((new_width, new_height), Image.LANCZOS)
        return resized_image

    def pad_or_crop(self, img, target_size, fill_color=(255, 255, 255)):
        iw, ih = img.size
        tw, th = target_size

        # 计算裁剪区域：若原图大于目标尺寸，则裁剪出中间部分；否则全部保留
        left = (iw - tw) // 2 if iw >= tw else 0
        top = (ih - th) // 2 if ih >= th else 0
        cropped = img.crop((left, top, left + min(iw, tw), top + min(ih, th)))

        # 新建目标尺寸的图像，并将裁剪后的图像居中粘贴
        new_img = Image.new(img.mode, target_size, fill_color)
        offset = ((tw - cropped.width) // 2, (th - cropped.height) // 2)
        new_img.paste(cropped, offset)

        return new_img
        
    def transform_foreground_original(self, original_fg, original_bg, transformation_info, canvas_size=400):
        drag_left = float(transformation_info.get("drag_left", 0))
        drag_top  = float(transformation_info.get("drag_top", 0))
        scale_ratio = float(transformation_info.get("scale_ratio", 1))
        data_orig_width  = float(transformation_info.get("data_original_width", canvas_size))
        data_orig_height = float(transformation_info.get("data_original_height", canvas_size))
        drag_width = float(transformation_info.get("drag_width", 0))
        drag_height = float(transformation_info.get("drag_height", 0))


        scale_ori_fg = canvas_size / max(original_fg.width, original_fg.height)
        scale_ori_bg = canvas_size / max(original_bg.width, original_bg.height)

        # 计算未缩放状态下（预览中）的默认居中位置（前景图未拖拽时的理想位置）
        default_left = (canvas_size - data_orig_width) / 2.0
        default_top  = (canvas_size - data_orig_height) / 2.0

        # 在未缩放状态下，计算实际拖拽产生的偏移（单位：像素，在预览尺寸下计算）
        offset_preview_x = drag_left - default_left
        offset_preview_y = drag_top - default_top

        offset_ori_x = offset_preview_x / scale_ori_fg
        offset_ori_y = offset_preview_y / scale_ori_fg

        new_width = int(original_fg.width * scale_ratio)
        new_height = int(original_fg.height * scale_ratio)
        scale_fg = original_fg.resize((new_width, new_height))

        output = Image.new("RGBA", (original_fg.width, original_fg.height), (255, 255, 255, 0))
        output.paste(scale_fg, (int(offset_ori_x), int(offset_ori_y)))

        new_width_fgbg = original_fg.width * scale_ori_fg / scale_ori_bg
        new_height_fgbg = original_fg.height * scale_ori_fg / scale_ori_bg
        scale_fgbg = output.resize((int(new_width_fgbg), int(new_height_fgbg)))


        final_output = Image.new("RGBA", (original_bg.width, original_bg.height), (255, 255, 255, 0))
        scale_fgbg = self.pad_or_crop(scale_fgbg, (original_bg.width, original_bg.height), (255, 255, 255, 0))
        final_output.paste(scale_fgbg, (0, 0))

        fit_fg = self.resize_to_fit_within(original_bg, original_fg)
        fit_fg = self.pad_or_crop(fit_fg, original_bg.size, (255, 255, 255, 0))

        return final_output, fit_fg

    @torch.inference_mode()
    def gradio_generate(self, background_img, foreground_img, transformation_info, seed, prompt, enable_gui, cfg=3.5, size_select="1024", text_strength=1, truecfg=False):
        try:
            trans = json.loads(transformation_info)
        except:
            trans = {}

        size_select = int(size_select)
        if size_select == 1024 and prompt != "": text_strength = 5
        if size_select == 768 and prompt != "": text_strength = 3
        
        r, g, b, ori_a = foreground_img.split()
        fg_img_scale, fg_img = self.transform_foreground_original(foreground_img, background_img, trans)

        new_r, new_g, new_b, new_a = fg_img_scale.split()
        foreground_img_scale = Image.merge("RGB", (new_r, new_g, new_b))

        r, g, b, ori_a = fg_img.split()
        foreground_img = Image.merge("RGB", (r, g, b))
        foreground_img_save = foreground_img.copy()
        ori_a = ori_a.convert("L")
        new_a = new_a.convert("L")
        foreground_img.paste((255, 255, 255), mask=ImageOps.invert(ori_a))

        images = self.model_generate(foreground_img.copy(), background_img.copy(),
                        ori_a, new_a,
                        enable_mask_affine=enable_gui, 
                        prompt=prompt,
                        offset_cond=[0, 1, 0] if not enable_gui else None,
                        seed=seed,
                        cfg=cfg,
                        size_select=size_select,
                        text_strength=text_strength,
                        truecfg=truecfg)
        images = Image.fromarray(images[0], "RGB")

        images = images.resize(background_img.size)
        # images.thumbnail((640, 640), Image.LANCZOS)
        return images


    @torch.inference_mode()
    def model_generate(self, fg_image, bg_image, ori_fg_mask, new_fg_mask, enable_mask_affine=True, prompt="", offset_cond=None, seed=None, cfg=3.5, size_select=1024, text_strength=1, truecfg=False):
        batch_size = 1
        # Prepare images
        # adjust bg->fg size
        fg_image, ori_fg_mask = adjust_fg_to_bg(fg_image, ori_fg_mask, bg_image.size)
        bucket_size = find_nearest_bucket_size(bg_image.size[0], bg_image.size[1], bucket_size=size_select)

        fg_image = fg_image.resize(bucket_size)
        bg_image = bg_image.resize(bucket_size)

        mask_affine = None
        if enable_mask_affine:
            ori_fg_mask = ori_fg_mask.resize(bucket_size)
            new_fg_mask = new_fg_mask.resize(bucket_size)
            mask_affine = get_mask_affine(new_fg_mask, ori_fg_mask)

        # Get embeddings
        prompt_embeds, pooled_prompt_embeds, text_ids = self._compute_text_embeddings(prompt)

        prompt_embeds = prompt_embeds.repeat(1, text_strength, 1)
        text_ids = text_ids.repeat(text_strength, 1)

        # Prepare 
        if self.config.model_choice == "dev":
            guidance = torch.full([1], cfg, device=self.device, dtype=torch.float32)
            guidance = guidance.expand(batch_size)
        else:
            guidance = None

        # Prepare generator
        if seed is None:
            seed = self.config.seed
        generator = torch.Generator(device=self.device).manual_seed(seed)
        # Prepare condition latents
        condition_image_latents = self._encode_images([fg_image, bg_image])
        
        if offset_cond is None:
            offset_cond = self.config.image_ids_offset
        offset_cond = offset_cond[1:]
        cond_latent_image_ids = []
        for offset_ in offset_cond:
            cond_latent_image_ids.append(
                self._prepare_image_ids(
                    condition_image_latents.shape[2] // 2,
                    condition_image_latents.shape[3] // 2,
                    offset_w=offset_ * condition_image_latents.shape[3] // 2
                )
            )

        if mask_affine is not None:
            affine_H, affine_W = condition_image_latents.shape[2] // 2, condition_image_latents.shape[3] // 2
            scale_factor = 1 / 16
            cond_latent_image_ids_fg = cond_latent_image_ids[0].reshape(affine_H, affine_W, 3).clone()

            # opt 1
            cond_latent_image_ids[0] = warp_affine_tensor(
                cond_latent_image_ids_fg, mask_affine, output_size=(affine_H, affine_W),
                scale_factor=scale_factor, device=self.device,
            )
        cond_latent_image_ids = torch.stack(cond_latent_image_ids)

        # Pack condition latents
        cond_image_latents = self._pack_latents(condition_image_latents)
        cond_input = {
            "image_latents": cond_image_latents,
            "image_ids": cond_latent_image_ids,
        }
        # Prepare initial latents
        width, height = bucket_size
        num_channels_latents = self.denoise_model.config.in_channels // 4
        latents, latent_image_ids = self._prepare_latents(
            batch_size, num_channels_latents, height, width, generator
        )

        # Setup timesteps
        sigmas = np.linspace(1.0, 1 / self.config.num_inference_steps, self.config.num_inference_steps)
        image_seq_len = latents.shape[1]
        mu = calculate_shift(
            image_seq_len,
            self.scheduler.config.base_image_seq_len,
            self.scheduler.config.max_image_seq_len,
            self.scheduler.config.base_shift,
            self.scheduler.config.max_shift,
        )
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler,
            self.config.num_inference_steps,
            self.device,
            sigmas=sigmas,
            mu=mu,
        )

        # Denoising loop
        for i, t in enumerate(timesteps):
            timestep = t.expand(latents.shape[0]).to(latents.dtype)
            with torch.autocast(enabled=True, device_type="cuda", dtype=self.config.dtype):
                noise_pred = self.denoise_model(
                    hidden_states=latents,
                    cond_input=cond_input,
                    timestep=timestep / 1000,
                    guidance=guidance,
                    pooled_projections=pooled_prompt_embeds,
                    encoder_hidden_states=prompt_embeds,
                    txt_ids=text_ids,
                    img_ids=latent_image_ids,
                    data_num_per_group=batch_size,
                    image_tags=self.config.image_tags,
                    context_tags=self.config.context_tags,
                    max_sequence_length=self.config.max_sequence_length,
                    mix_attention_double=self.config.mix_attention_double,
                    mix_attention_single=self.config.mix_attention_single,
                    joint_attention_kwargs=None,
                    return_dict=False,
                )[0]

                if truecfg and i >= 1:
                    guidance_neg = torch.full([1], 1, device=self.device, dtype=torch.float32)
                    guidance_neg = guidance_neg.expand(batch_size)
                    noise_pred_neg = self.denoise_model(
                        hidden_states=latents,
                        cond_input=cond_input,
                        timestep=timestep / 1000,
                        guidance=guidance,
                        pooled_projections=pooled_prompt_embeds,
                        encoder_hidden_states=prompt_embeds,
                        txt_ids=text_ids,
                        img_ids=latent_image_ids,
                        data_num_per_group=batch_size,
                        image_tags=self.config.image_tags,
                        context_tags=self.config.context_tags,
                        max_sequence_length=self.config.max_sequence_length,
                        mix_attention_double=self.config.mix_attention_double,
                        mix_attention_single=self.config.mix_attention_single,
                        joint_attention_kwargs=None,
                        return_dict=False,
                    )[0]
                    noise_pred = noise_pred_neg + 5 * (noise_pred - noise_pred_neg)

            # Compute previous noisy sample
            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]

        # Decode latents
        latents = self._unpack_latents(latents, height, width)
        latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
        images = self.vae.decode(latents, return_dict=False)[0]

        # Post-process images
        images = images.add(1).mul(127.5).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
        return images

    def _encode_images(self, images):
        return encode_images_cond(self.vae, [images], self.device)

    def _prepare_image_ids(self, h, w, offset_w=0):
        return _prepare_image_ids(h, w, offset_w=offset_w).to(self.device)

    def _pack_latents(self, latents):
        b, c, h, w = latents.shape
        return _pack_latents(latents, b, c, h, w)

    def _unpack_latents(self, latents, height, width):
        vae_scale = 2 ** (len(self.vae.config.block_out_channels) - 1)
        return _unpack_latents(latents, height, width, vae_scale)

    def _prepare_latents(self, batch_size, num_channels_latents, height, width, generator):
        vae_scale = 2 ** (len(self.vae.config.block_out_channels) - 1)
        latents, latent_image_ids = prepare_latents(
            batch_size=batch_size,
            num_channels_latents=num_channels_latents,
            vae_downsample_factor=vae_scale,
            height=height,
            width=width,
            dtype=self.config.dtype,
            device=self.device,
            generator=generator,
            offset=None
        )
        return latents, latent_image_ids