visualcloze.py

import random
from einops import rearrange
from diffusers.models import AutoencoderKL
from PIL import Image
import torch
import torch.nn.functional as F
from torchvision import transforms
from torchvision.transforms.functional import to_pil_image
from models.sampling import prepare_modified
from models.util import load_clip, load_t5, load_flow_model
from transport import Sampler, create_transport
from data.imgproc import to_rgb_if_rgba


def center_crop(image, target_size):
    width, height = image.size
    new_width, new_height = target_size

    left = (width - new_width) // 2
    top = (height - new_height) // 2
    right = left + new_width
    bottom = top + new_height

    return image.crop((left, top, right, bottom))


def resize_with_aspect_ratio(img, resolution, divisible=16, aspect_ratio=None):
    """Resize image while maintaining aspect ratio, ensuring area is close to resolution**2 and dimensions are divisible by 16
    
    Args:
        img: PIL Image or torch.Tensor (C,H,W)/(B,C,H,W)
        resolution: target resolution
        divisible: ensure output dimensions are divisible by this number
    
    Returns:
        Resized image of the same type as input
    """
    # Check input type and get dimensions
    is_tensor = isinstance(img, torch.Tensor)
    if is_tensor:
        if img.dim() == 3:
            c, h, w = img.shape
            batch_dim = False
        else:
            b, c, h, w = img.shape
            batch_dim = True
    else:
        w, h = img.size
        
    # Calculate new dimensions
    if aspect_ratio is None:
        aspect_ratio = w / h
    target_area = resolution * resolution
    new_h = int((target_area / aspect_ratio) ** 0.5)
    new_w = int(new_h * aspect_ratio)
    
    # Ensure divisible by divisible
    new_w = max(new_w // divisible, 1) * divisible
    new_h = max(new_h // divisible, 1) * divisible
    
    # Adjust size based on input type
    if is_tensor:
        # Use torch interpolation method
        mode = 'bilinear'
        align_corners = False
        if batch_dim:
            return F.interpolate(img, size=(new_h, new_w), 
                               mode=mode, align_corners=align_corners)
        else:
            return F.interpolate(img.unsqueeze(0), size=(new_h, new_w),
                               mode=mode, align_corners=align_corners).squeeze(0)
    else:
        # Use PIL LANCZOS resampling
        return img.resize((new_w, new_h), Image.LANCZOS)


class VisualClozeModel:
    def __init__(
        self, model_path, model_name="flux-dev-fill-lora", max_length=512, lora_rank=256, 
        atol=1e-6, rtol=1e-3, solver='euler', time_shifting_factor=1, 
        resolution=384, precision='bf16'):
        self.atol = atol
        self.rtol = rtol
        self.solver = solver
        self.time_shifting_factor = time_shifting_factor
        self.resolution = resolution
        self.precision = precision
        self.max_length = max_length
        self.lora_rank = lora_rank
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[self.precision]
        
        # Initialize model
        print("Initializing model...")
        self.model = load_flow_model(model_name, device=self.device, lora_rank=self.lora_rank)
        
        # Initialize VAE
        print("Initializing VAE...")
        self.ae = AutoencoderKL.from_pretrained(f"black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=self.dtype).to(self.device)
        self.ae.requires_grad_(False)
        
        # Initialize text encoders
        print("Initializing text encoders...")
        self.t5 = load_t5(self.device, max_length=self.max_length)
        self.clip = load_clip(self.device)
        
        self.model.eval().to(self.device, dtype=self.dtype)
        
        # Load model weights
        ckpt = torch.load(model_path)
        self.model.load_state_dict(ckpt, strict=False)
        del ckpt
        
        # Initialize sampler
        transport = create_transport(
            "Linear",
            "velocity",
            do_shift=True,
        ) 
        self.sampler = Sampler(transport)
        self.sample_fn = self.sampler.sample_ode(
            sampling_method=self.solver,
            num_steps=30,
            atol=self.atol,
            rtol=self.rtol,
            reverse=False,
            do_shift=True,
            time_shifting_factor=self.time_shifting_factor,
        )
        
        # Image transformation
        self.image_transform = transforms.Compose([
            transforms.Lambda(lambda img: to_rgb_if_rgba(img)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])
        
        self.grid_h = None
        self.grid_w = None
        
    def set_grid_size(self, h, w):
        """Set grid size"""
        self.grid_h = h
        self.grid_w = w
    
    @torch.no_grad
    def upsampling(self, image, target_size, cfg, upsampling_steps, upsampling_noise, generator, content_prompt):
        content_instruction = [
            "The content of the last image in the final row is: ",
            "The last image of the last row depicts: ",
            "In the final row, the last image shows: ",
            "The last image in the bottom row illustrates: ",
            "The content of the bottom-right image is: ",
            "The final image in the last row portrays: ",
            "The last image of the final row displays: ",
            "In the last row, the final image captures: ",
            "The bottom-right corner image presents: ",
            "The content of the last image in the concluding row is: ",
            "In the last row, ",
            "The editing instruction in the last row is: ", 
        ]
        for c in content_instruction:
            if content_prompt.startswith(c):
                content_prompt = content_prompt.replace(c, '')
        
        if target_size is None:
            aspect_ratio = 1
            target_area = 1024 * 1024
            new_h = int((target_area / aspect_ratio) ** 0.5)
            new_w = int(new_h * aspect_ratio)
            target_size = (new_w, new_h)

        if target_size[0] * target_size[1] > 1024 * 1024:
            aspect_ratio = target_size[0] / target_size[1]
            target_area = 1024 * 1024
            new_h = int((target_area / aspect_ratio) ** 0.5)
            new_w = int(new_h * aspect_ratio)
            target_size = (new_w, new_h)
        
        image = image.resize(((target_size[0] // 16) * 16, (target_size[1] // 16) * 16))
        if upsampling_noise >= 1.0:
            return image

        self.sample_fn = self.sampler.sample_ode(
            sampling_method=self.solver,
            num_steps=upsampling_steps,
            atol=self.atol,
            rtol=self.rtol,
            reverse=False,
            do_shift=False,
            time_shifting_factor=1.0, 
            strength=upsampling_noise
        )

        processed_image = self.image_transform(image)
        processed_image = processed_image.to(self.device, non_blocking=True)
        blank = torch.zeros_like(processed_image, device=self.device, dtype=self.dtype)
        mask = torch.full((1, 1, processed_image.shape[1], processed_image.shape[2]), fill_value=1, device=self.device, dtype=self.dtype)
        with torch.no_grad():
            latent = self.ae.encode(processed_image[None].to(self.ae.dtype)).latent_dist.sample()
            blank = self.ae.encode(blank[None].to(self.ae.dtype)).latent_dist.sample()
            latent = (latent - self.ae.config.shift_factor) * self.ae.config.scaling_factor
            blank = (blank - self.ae.config.shift_factor) * self.ae.config.scaling_factor
            latent_h, latent_w = latent.shape[2:]

            mask = rearrange(mask, "b c (h ph) (w pw) -> b (c ph pw) h w", ph=8, pw=8) 
            mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
            
            latent = latent.to(self.dtype)
            blank = blank.to(self.dtype)
            latent = rearrange(latent, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
            blank = rearrange(blank, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
            
            img_cond = torch.cat((blank, mask), dim=-1)
    
            # Generate noise
            noise = torch.randn([1, 16, latent_h, latent_w], device=self.device, generator=generator).to(self.dtype)
            x = [[noise]]
            
            inp = prepare_modified(t5=self.t5, clip=self.clip, img=x, prompt=[content_prompt], proportion_empty_prompts=0.0)
            inp["img"] = inp["img"] * (1 - upsampling_noise) + latent * upsampling_noise
            model_kwargs = dict(
                txt=inp["txt"], 
                txt_ids=inp["txt_ids"], 
                txt_mask=inp["txt_mask"],
                y=inp["vec"], 
                img_ids=inp["img_ids"], 
                img_mask=inp["img_mask"], 
                cond=img_cond,
                guidance=torch.full((1,), cfg, device=self.device, dtype=self.dtype),
            )
            sample = self.sample_fn(
                inp["img"], self.model.forward, model_kwargs
            )[-1]
            
            sample = sample[:1]
            sample = rearrange(sample, "b (h w) (c ph pw) -> b c (h ph) (w pw)", ph=2, pw=2, h=latent_h // 2, w=latent_w // 2)
            sample = self.ae.decode(sample / self.ae.config.scaling_factor + self.ae.config.shift_factor)[0]
            sample = (sample + 1.0) / 2.0
            sample.clamp_(0.0, 1.0)
            sample = sample[0]
            
            output_image = to_pil_image(sample.float())
            
            return output_image
    
    def process_images(
        self, images: list[list[Image.Image]], 
        prompts: list[str], 
        seed: int = 0, 
        cfg: int = 30, 
        steps: int = 30, 
        upsampling_steps: int = 10, 
        upsampling_noise: float = 0.4, 
        is_upsampling: bool =True):
        """
        Processes a list of images based on the provided text prompts and settings, with optional upsampling to enhance image resolution or detail.

        Parameters:
            images (list[list[Image.Image]]): A collection of images arranged in a grid layout, where each row represents an in-context example or the current query. 
            The current query should be placed in the last row. The target image may be None in the input, while all other images should be of the PIL Image type (Image.Image).
            
            prompts (list[str]): A list containing three prompts: the layout prompt, task prompt, and content prompt, respectively.
            
            seed (int): A fixed integer seed to ensure reproducibility of random elements during processing.
            
            cfg (int): The strength of Classifier-Free Diffusion Guidance, which controls the degree of influence over the generated results.
            
            steps (int): The number of sampling steps to be performed during processing.
            
            upsampling_steps (int): The number of denoising steps to apply when performing upsampling.
            
            upsampling_noise (float): The noise level used as a starting point when upsampling with SDEdit. A higher value reduces noise, and setting it to 1 disables SDEdit, causing the PIL resize function to be used instead.
            
            is_upsampling (bool, optional): A flag indicating whether upsampling should be applied using SDEdit.

        Returns:
            Processed images resulting from the algorithm, with optional upsampling applied based on the `is_upsampling` flag.
        """
        
        if seed == 0:
            seed = random.randint(0, 2 ** 32 - 1)
        
        self.sample_fn = self.sampler.sample_ode(
            sampling_method=self.solver,
            num_steps=steps,
            atol=self.atol,
            rtol=self.rtol,
            reverse=False,
            do_shift=True,
            time_shifting_factor=self.time_shifting_factor,
        )

        # Use class grid size
        grid_h, grid_w = self.grid_h, self.grid_w
        
        # Ensure all images are RGB mode or None
        for i in range(0, grid_h):
            images[i] = [img.convert("RGB") if img is not None else None for img in images[i]]
        
        # Adjust all image sizes
        resolution = self.resolution
        processed_images = []
        mask_position = []
        target_size = None
        upsampling_size = None
        
        for i in range(grid_h):
            # Find the size of the first non-empty image in this row
            reference_size = None
            for j in range(0, grid_w):
                if images[i][j] is not None:
                    if i == grid_h - 1 and upsampling_size is None:
                        upsampling_size = images[i][j].size

                    resized = resize_with_aspect_ratio(images[i][j], resolution, aspect_ratio=None)
                    reference_size = resized.size
                    if i == grid_h - 1 and target_size is None:
                        target_size = reference_size
                    break
            
            # Process all images in this row
            for j in range(0, grid_w):
                if images[i][j] is not None:
                    target = resize_with_aspect_ratio(images[i][j], resolution, aspect_ratio=None)
                    if target.width <= target.height:
                        target = target.resize((reference_size[0], int(reference_size[0] / target.width * target.height)))
                        target = center_crop(target, reference_size)
                    elif target.width > target.height:
                        target = target.resize((int(reference_size[1] / target.height * target.width), reference_size[1]))
                        target = center_crop(target, reference_size)
                    
                    processed_images.append(target)
                    if i == grid_h - 1:
                        mask_position.append(0)
                else:
                    # If this row has a reference size, use it; otherwise use default size
                    if reference_size:
                        blank = Image.new('RGB', reference_size, (0, 0, 0))
                    else:
                        blank = Image.new('RGB', (resolution, resolution), (0, 0, 0))
                    processed_images.append(blank)
                    if i == grid_h - 1:
                        mask_position.append(1)
                    else:
                        raise ValueError('Please provide each image in the in-context example.')
            
        # return processed_images
        
        if len(mask_position) > 1 and sum(mask_position) > 1:
            if target_size is None:
                new_w = 384
            else:
                new_w = target_size[0]
            for i in range(len(processed_images)):
                if processed_images[i] is not None:
                    new_h = int(processed_images[i].height * (new_w / processed_images[i].width))
                    new_w = int(new_w / 16) * 16
                    new_h = int(new_h / 16) * 16
                    processed_images[i] = processed_images[i].resize((new_w, new_h))
                
        # Build grid image and mask
        with torch.autocast("cuda", self.dtype):
            grid_image = []
            fill_mask = []
            for i in range(grid_h):
                row_images = [self.image_transform(img) for img in processed_images[i * grid_w: (i + 1) * grid_w]]
                if i == grid_h - 1:
                    row_masks = [torch.full((1, 1, row_images[0].shape[1], row_images[0].shape[2]), fill_value=m, device=self.device) for m in mask_position]
                else:
                    row_masks = [torch.full((1, 1, row_images[0].shape[1], row_images[0].shape[2]), fill_value=0, device=self.device) for m in mask_position]

                grid_image.append(torch.cat(row_images, dim=2).to(self.device, non_blocking=True))
                fill_mask.append(torch.cat(row_masks, dim=3))
            # Encode condition image
            with torch.no_grad():
                fill_cond = [self.ae.encode(img[None].to(self.ae.dtype)).latent_dist.sample()[0] for img in grid_image]
                fill_cond = [(img - self.ae.config.shift_factor) * self.ae.config.scaling_factor for img in fill_cond]
                
                # Rearrange mask
                fill_mask = [rearrange(mask, "b c (h ph) (w pw) -> b (c ph pw) h w", ph=8, pw=8) for mask in fill_mask]
                fill_mask = [rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) for mask in fill_mask]
            
            fill_cond = [img.to(self.dtype) for img in fill_cond]
            fill_cond = [rearrange(img.unsqueeze(0), "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) for img in fill_cond]
            
            fill_cond =  torch.cat(fill_cond, dim=1)
            fill_mask =  torch.cat(fill_mask, dim=1)
            img_cond = torch.cat((fill_cond, fill_mask), dim=-1)
        
            # Generate sample
            noise = []
            sliced_subimage = []
            rng = torch.Generator(device=self.device).manual_seed(int(seed))
            for sub_img in grid_image:
                h, w = sub_img.shape[-2:]
                sliced_subimage.append((h, w))
                latent_w, latent_h = w // 8, h // 8
                noise.append(torch.randn([1, 16, latent_h, latent_w], device=self.device, generator=rng).to(self.dtype))
            x = [noise]
            
            with torch.no_grad():
                inp = prepare_modified(t5=self.t5, clip=self.clip, img=x, prompt=[' '.join(prompts)], proportion_empty_prompts=0.0)
                
                model_kwargs = dict(
                    txt=inp["txt"], 
                    txt_ids=inp["txt_ids"], 
                    txt_mask=inp["txt_mask"],
                    y=inp["vec"], 
                    img_ids=inp["img_ids"], 
                    img_mask=inp["img_mask"], 
                    cond=img_cond,
                    guidance=torch.full((1,), cfg, device=self.device, dtype=self.dtype),
                )
                samples = self.sample_fn(
                    inp["img"], self.model.forward, model_kwargs
                )[-1]

            # Get query row
            with torch.no_grad():
                samples = samples[:1]
                row_samples = []
                start = 0
                for size in sliced_subimage:
                    end = start + (size[0] * size[1] // 256)
                    latent_h = size[0] // 8
                    latent_w = size[1] // 8
                    row_sample = samples[:, start:end, :]
                    row_sample = rearrange(row_sample, "b (h w) (c ph pw) -> b c (h ph) (w pw)", ph=2, pw=2, h=latent_h//2, w=latent_w//2)
                    row_sample = self.ae.decode(row_sample / self.ae.config.scaling_factor + self.ae.config.shift_factor)[0]
                    row_sample = (row_sample + 1.0) / 2.0
                    row_sample.clamp_(0.0, 1.0)
                    row_samples.append(row_sample[0])
                    start = end
            
            # Convert all samples to PIL images
            output_images = []
            for row_sample in row_samples:
                output_image = to_pil_image(row_sample.float())
                output_images.append(output_image)
            
            torch.cuda.empty_cache()
            
            ret = []
            ret_w = output_images[-1].width
            ret_h = output_images[-1].height
            
            row_start = (grid_h - 1) * grid_w
            row_end = grid_h * grid_w
            for i in range(row_start, row_end):
                # when the image is masked, then output it
                if mask_position[i - row_start] and is_upsampling:
                    cropped = output_images[-1].crop(((i - row_start) * ret_w // self.grid_w, 0, ((i - row_start) + 1) * ret_w // self.grid_w, ret_h))
                    upsampled = self.upsampling(
                        cropped, 
                        upsampling_size, 
                        cfg, 
                        upsampling_steps=upsampling_steps, 
                        upsampling_noise=upsampling_noise, 
                        generator=rng, 
                        content_prompt=prompts[2])
                    ret.append(upsampled)
                elif mask_position[i - row_start]:
                    cropped = output_images[-1].crop(((i - row_start) * ret_w // self.grid_w, 0, ((i - row_start) + 1) * ret_w // self.grid_w, ret_h))
                    ret.append(cropped)
            
            return ret