inference.py

# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import dataclasses
from typing import Literal

from accelerate import Accelerator
from transformers import HfArgumentParser
from PIL import Image
import json
import itertools

from uno.flux.pipeline import UNOPipeline, preprocess_ref
from uno.utils.image_utils import concat_images, random_bbox, draw_rectangle, bbox_to_mask, center_crop_with_bbox, analyze_mask_overlap_gpu
from torch.utils.data import DataLoader
import torch
def horizontal_concat(images):
    widths, heights = zip(*(img.size for img in images))

    total_width = sum(widths)
    max_height = max(heights)

    new_im = Image.new('RGB', (total_width, max_height))

    x_offset = 0
    for img in images:
        new_im.paste(img, (x_offset, 0))
        x_offset += img.size[0]

    return new_im

@dataclasses.dataclass
class InferenceArgs:
    prompt: str | None = None
    image_paths: list[str] | None = None
    eval_json_path: str | None = None
    offload: bool = False
    num_images_per_prompt: int = 1
    model_type: Literal["flux-dev", "flux-dev-fp8", "flux-schnell"] = "flux-dev"
    width: int = 512
    height: int = 512
    ref_size: int = -1
    num_steps: int = 25
    guidance: float = 4
    seed: int = 3407
    mask_type: str = None

    rope_type: str = "uno"  #uno/flux   2d/3d rope
    saved_dir: str = "None"
    dit_lora_path: str = "None"

    #overlap#
    use_nerf: bool = True
    a: int = 0
    #########

    only_lora: bool = True
    concat_refs: bool = False
    lora_rank: int = 512
    pe: Literal['d', 'h', 'w', 'o'] = 'd'
  
def main(args: InferenceArgs):
    accelerator = Accelerator()

    #overlap
    args.overlap_config = {
        "use_nerf": args.use_nerf,
        "a": args.a,
    }
    #
    pipeline = UNOPipeline(
        args.model_type,
        accelerator.device,
        args.offload,
        only_lora=args.only_lora,
        lora_rank=args.lora_rank,
        dit_lora_path = args.dit_lora_path,
        args = args,
    )

    assert args.prompt is not None or args.eval_json_path is not None, \
        "Please provide either prompt or eval_json_path"
    
    if args.eval_json_path is not None:
        with open(args.eval_json_path, "rt") as f:
            data_dicts = json.load(f)
        data_root = os.path.dirname(args.eval_json_path)
    else:
        data_root = "./"
        data_dicts = [{"prompt": args.prompt, "image_paths": args.image_paths}]

    for (i, data_dict), j in itertools.product(enumerate(data_dicts), range(args.num_images_per_prompt)):
        if (i * args.num_images_per_prompt + j) % accelerator.num_processes != accelerator.process_index:
            continue

        ref_imgs = [
            Image.open(img_path)
            for img_path in (data_dict["ref_img"] if isinstance(data_dict["ref_img"],list) else [data_dict["ref_img"]])
        ]
        if args.ref_size==-1:
            args.ref_size = 512

        #create bbox
        img_bbox_ = data_dict["img_bbox"]    #[[],[],[]]
        img_bbox = []
        USE_BBOX_NORM = False 
        for one_bbox in img_bbox_:
            if args.mask_type =="none":
                img_bbox.append([0, 0, args.width, args.height])
                continue
            if one_bbox[0]<=1.0:
                mid_x = (one_bbox[0] + one_bbox[2])/2 * args.width
                mid_y = (one_bbox[1] + one_bbox[3])/2 * args.height
                bbox_width = (one_bbox[2] - one_bbox[0]) 
                bbox_height = (one_bbox[3] - one_bbox[1])
                if USE_BBOX_NORM:
                    ratio =  bbox_height / bbox_width
                    if args.height <= args.width:
                        bbox_height = bbox_height * args.height
                        bbox_width = bbox_height/ratio
                    else:
                        bbox_width = bbox_width * args.width
                        bbox_height = bbox_width * ratio
                else:
                    bbox_height = bbox_height * args.height
                    bbox_width = bbox_width * args.width
                left = int(mid_x - bbox_width/2)
                right = int(mid_x + bbox_width/2)
                up = int(mid_y - bbox_height/2)
                down = int(mid_y + bbox_height/2)
                img_bbox.append([left,up,right,down])
                #img_bbox.append([one_bbox[0] * args.width, one_bbox[1] * args.height, one_bbox[2] * args.width, one_bbox[3] * args.height])
            else:
                img_bbox.append(one_bbox)

        ref_imgs_new = []
        for index, img in enumerate(ref_imgs):
            ref_img_new, _ = center_crop_with_bbox(img, None, args.ref_size, args.ref_size) 
            ref_imgs_new.append(ref_img_new)


        img_bbox_mask = torch.stack([ torch.stack([bbox_to_mask(bbox, args.height, args.width, resize = True) for bbox in one_batch])  for one_batch in [img_bbox]])

        #overlap
        overlap_info, fg_count = analyze_mask_overlap_gpu(img_bbox_mask, img_bbox_mask.device)

        prompt = data_dict["prompt"]
        print("Prompt:",prompt)
        image_gen, pure_img = pipeline(
            prompt=prompt,
            width=args.width,
            height=args.height,
            guidance=args.guidance,
            num_steps=args.num_steps,
            seed=args.seed,
            ref_imgs=ref_imgs_new,
            pe=args.pe,
            img_bboxes = img_bbox,
            target_mask = img_bbox_mask,
            rope_type = args.rope_type,

            #overlap
            overlap={"overlap_info":overlap_info, "fg_count":fg_count},
        )
        os.makedirs(args.save_path, exist_ok=True)
        image_gen.save(os.path.join(args.save_path, f"{i}_{j}.png"))

        os.makedirs(args.save_path + '_pure', exist_ok=True)
        pure_img.save(os.path.join(args.save_path + '_pure', f"{i}_{j}.png"))

        # save config and image
        args_dict = vars(args)
        args_dict['prompt'] = data_dict["prompt"]
        args_dict['image_paths'] = data_dict["ref_img"]
        #with open(os.path.join(args.save_path, f"{i}_{j}.json"), 'w') as f:
        #    json.dump(args_dict, f, indent=4)        

if __name__ == "__main__":
    parser = HfArgumentParser([InferenceArgs])
    args = parser.parse_args_into_dataclasses()[0]

    from datetime import datetime
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    args.save_path = os.path.join(args.saved_dir,f"res_seed{args.seed}_w{args.width}h{args.height}_ref{args.ref_size}_alpha-{args.a}_{timestamp}")
    main(args)