infer.py

"""
Image inference module for VAREdit model.
Supports 2B and 8B model variants for image editing with text instructions.
"""
import argparse
import logging
from typing import Tuple, Any, Optional
from torchvision.transforms.functional import to_tensor
import numpy as np
from PIL import Image
import PIL.Image as PImage
from tools.run_infinity import (
    load_tokenizer, load_visual_tokenizer, load_transformer,
    gen_one_img, h_div_w_templates, dynamic_resolution_h_w
)
import time
import torch

def transform(pil_img, target_image_size):
    # currently only support square image.
    width, height = pil_img.size
    max_dim = max(width, height)
    padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
    padded_image.paste(pil_img, (0, 0))
    def crop_op(image):
        image = image.resize((max_dim, max_dim), resample=PImage.LANCZOS)
        crop_image = image.crop((0, 0, width, height))
        return crop_image
    padded_image = padded_image.resize((target_image_size, target_image_size), resample=PImage.LANCZOS)
    im = to_tensor(np.array(padded_image))
    return im.add(im).add_(-1), crop_op

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Model configurations
MODEL_CONFIGS = {
    '2B': {
        'vae_filename': 'infinity_vae_d32reg.pth',
        'vae_type': 32,
        'model_type': 'infinity_2b',
        'apply_spatial_patchify': 0,
    },
    '8B': {
        'vae_filename': 'infinity_vae_d56_f8_14_patchify.pth',
        'vae_type': 14,
        'model_type': 'infinity_8b',
        'apply_spatial_patchify': 1,
    }
}

# Common model arguments
COMMON_ARGS = {
    'cfg_insertion_layer': 0,
    'add_lvl_embeding_only_first_block': 1,
    'use_bit_label': 1,
    'rope2d_each_sa_layer': 1,
    'rope2d_normalized_by_hw': 2,
    'use_scale_schedule_embedding': 0,
    'sampling_per_bits': 1,
    'text_channels': 2048,
    'h_div_w_template': 1.000,
    'use_flex_attn': 0,
    'cache_dir': '/dev/shm',
    'checkpoint_type': 'torch',
    'bf16': 1,
    'enable_model_cache': 0,
}


def load_model(pretrain_root: str, model_path: str, model_size: str, image_size: int) -> Tuple[Any, ...]:
    """
    Load the model and its components.
    
    Args:
        pretrain_root: Root directory for pretrained models
        model_path: Path to the specific model checkpoint
        
    Returns:
        Tuple of (args, model, vae, tokenizer, text_encoder)
        
    Raises:
        ValueError: If unsupported model size is specified
    """
    if model_size not in MODEL_CONFIGS:
        raise ValueError(f"Unsupported model size: {model_size}. Choose '2B' or '8B'.")
    
    config = MODEL_CONFIGS[model_size]
    
    # Build arguments
    args_dict = {
        **COMMON_ARGS,
        **config,
        'model_path': model_path,
        'vae_path': f"{pretrain_root}/{config['vae_filename']}",
        'text_encoder_ckpt': f"{pretrain_root}/flan-t5-xl"
    }
    args = argparse.Namespace(**args_dict)
    if image_size == 512:
        args.pn = "0.25M"
    elif image_size == 1024:
        args.pn = "1M"
    else:
        raise ValueError(f"Unsupported image size: {image_size}. Choose 512 or 1024.")
    logger.info(f"Loading {model_size} model from {model_path}")
    
    # Load components
    text_tokenizer, text_encoder = load_tokenizer(t5_path=args.text_encoder_ckpt)
    vae = load_visual_tokenizer(args)
    model = load_transformer(vae, args)
    
    logger.info("Model loaded successfully")
    return args, model, vae, text_tokenizer, text_encoder, image_size


def generate_image(
    model_components: Tuple[Any, ...],
    src_img_path: str,
    instruction: str,
    cfg: float = 4.0,
    tau: float = 0.5,
    seed: Optional[int] = -1,
) -> None:
    """
    Generate edited image based on source image and text instruction.
    
    Args:
        model_components: Tuple of (args, model, vae, tokenizer, text_encoder)
        src_img_path: Path to source image
        instruction: Text instruction for editing
        cfg: Classifier-free guidance scale
        tau: Temperature parameter
    """
    args, model, vae, tokenizer, text_encoder, image_size = model_components
    
    # Set default image size
    assert image_size in [512, 1024], f"Invalid image size: {image_size}, expected 512 or 1024"
    if image_size == 512:
        pn = "0.25M"
    elif image_size == 1024:
        pn = "1M"
    
    # Load and preprocess source image
    try:
        with Image.open(src_img_path) as src_img:
            src_img = src_img.convert('RGB')
            src_img_tensor, crop_op = transform(src_img, image_size)
    except Exception as e:
        logger.error(f"Failed to load source image: {e}")
        raise
    
    # Set up generation parameters
    aspect_ratio = 1.0  # h:w ratio
    h_div_w_template = h_div_w_templates[np.argmin(np.abs(h_div_w_templates - aspect_ratio))]
    scale_schedule = [(1, h, w) for (_, h, w) in dynamic_resolution_h_w[h_div_w_template][pn]['scales']]
    
    logger.info(f"Generating image with instruction: '{instruction}'")
    
    # Generate image
    if seed == -1:
        seed = np.random.randint(0, 1000000)
    torch.cuda.empty_cache()
    start_time = time.time()
    generated_image = gen_one_img(
        model, vae, tokenizer, text_encoder,
        instruction, src_img_tensor,
        g_seed=seed,
        gt_leak=0,
        gt_ls_Bl=None,
        cfg_list=cfg,
        tau_list=tau,
        scale_schedule=scale_schedule,
        cfg_insertion_layer=[args.cfg_insertion_layer],
        vae_type=args.vae_type,
        sampling_per_bits=args.sampling_per_bits,
        enable_positive_prompt=0,
        apply_spatial_patchify=args.apply_spatial_patchify,
    )
    end_time = time.time()
    logger.info(f"Time taken: {end_time - start_time:.2f} seconds")
    max_memory = torch.cuda.max_memory_allocated() / 1024 ** 3
    logger.info(f"Max memory: {max_memory:.2f} GB")
    generated_image_np = generated_image.cpu().numpy()
    if generated_image_np.shape[2] == 3:
        generated_image_np = generated_image_np[..., ::-1]
    result_image = Image.fromarray(generated_image_np.astype(np.uint8))
    result_image = crop_op(result_image)
    return result_image

def main():
    """Main execution function with example usage."""
    try:
        # Load model
        model_components = load_model(
            "HiDream-ai/VAREdit",
            "HiDream-ai/VAREdit/8B-1024.pth",
            "8B",
            1024
        )
        
        # Generate image
        generate_image(
            model_components,
            "assets/test.jpg",
            "Add glasses to this girl and change hair color to red",
            cfg=3.0,
            tau=1.0,
            seed=42
        )

    except Exception as e:
        logger.error(f"Inference failed: {e}")
        raise


if __name__ == "__main__":
    main()