197 lines
6.6 KiB
Python
197 lines
6.6 KiB
Python
# 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
|
|
import torch
|
|
|
|
from uso.flux.pipeline import USOPipeline, preprocess_ref
|
|
from transformers import SiglipVisionModel, SiglipImageProcessor
|
|
from tqdm import tqdm
|
|
|
|
|
|
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
|
|
offload: bool = True
|
|
num_images_per_prompt: int = 1
|
|
model_type: Literal["flux-dev", "flux-dev-fp8", "flux-schnell"] = "flux-dev-fp8"
|
|
width: int = 1024
|
|
height: int = 1024
|
|
num_steps: int = 25
|
|
guidance: float = 4
|
|
seed: int = 3407
|
|
save_path: str = "output/inference"
|
|
only_lora: bool = True
|
|
concat_refs: bool = False
|
|
lora_rank: int = 128
|
|
pe: Literal["d", "h", "w", "o"] = "d"
|
|
content_ref: int = 512
|
|
ckpt_path: str | None = None
|
|
use_siglip: bool = True
|
|
instruct_edit: bool = False
|
|
hf_download: bool = True
|
|
|
|
|
|
def main(args: InferenceArgs):
|
|
accelerator = Accelerator()
|
|
|
|
# init SigLIP model
|
|
siglip_processor = None
|
|
siglip_model = None
|
|
|
|
siglip_path = '/workspace/comfyui/models/clip/siglip-so400m-patch14-384'
|
|
if args.use_siglip:
|
|
siglip_processor = SiglipImageProcessor.from_pretrained(
|
|
# "google/siglip-so400m-patch14-384"
|
|
siglip_path
|
|
)
|
|
siglip_model = SiglipVisionModel.from_pretrained(
|
|
# "google/siglip-so400m-patch14-384"
|
|
siglip_path
|
|
)
|
|
siglip_model.eval()
|
|
siglip_model.to(accelerator.device)
|
|
print("SigLIP model loaded successfully")
|
|
|
|
pipeline = USOPipeline(
|
|
args.model_type,
|
|
accelerator.device,
|
|
args.offload,
|
|
only_lora=args.only_lora,
|
|
lora_rank=args.lora_rank,
|
|
hf_download=args.hf_download,
|
|
)
|
|
if args.use_siglip and siglip_model is not None:
|
|
pipeline.model.vision_encoder = siglip_model
|
|
print('-----> hook siglip encoder')
|
|
|
|
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}]
|
|
|
|
print(
|
|
f"process: {accelerator.num_processes}/{accelerator.process_index}, \
|
|
process images: {len(data_dicts)}/{len(data_dicts[accelerator.process_index::accelerator.num_processes])}"
|
|
)
|
|
|
|
data_dicts = data_dicts[accelerator.process_index :: accelerator.num_processes]
|
|
|
|
accelerator.wait_for_everyone()
|
|
local_task_count = len(data_dicts) * args.num_images_per_prompt
|
|
if accelerator.is_main_process:
|
|
progress_bar = tqdm(total=local_task_count, desc="Generating Images")
|
|
|
|
for (i, data_dict), j in itertools.product(
|
|
enumerate(data_dicts), range(args.num_images_per_prompt)
|
|
):
|
|
ref_imgs = []
|
|
for _, img_path in enumerate(data_dict["image_paths"]):
|
|
if img_path != "":
|
|
img = Image.open(os.path.join(data_root, img_path)).convert("RGB")
|
|
ref_imgs.append(img)
|
|
else:
|
|
ref_imgs.append(None)
|
|
siglip_inputs = None
|
|
if args.use_siglip and siglip_processor is not None:
|
|
with torch.no_grad():
|
|
siglip_inputs = [
|
|
siglip_processor(img, return_tensors="pt").to(pipeline.device)
|
|
for img in ref_imgs[1:] if isinstance(img, Image.Image)
|
|
]
|
|
|
|
ref_imgs_pil = [
|
|
preprocess_ref(img, args.content_ref) for img in ref_imgs[:1] if isinstance(img, Image.Image)
|
|
]
|
|
|
|
if args.instruct_edit:
|
|
args.width, args.height = ref_imgs_pil[0].size
|
|
args.width, args.height = args.width * (1024 / args.content_ref), args.height * (1024 / args.content_ref)
|
|
image_gen = pipeline(
|
|
prompt=data_dict["prompt"],
|
|
width=args.width,
|
|
height=args.height,
|
|
guidance=args.guidance,
|
|
num_steps=args.num_steps,
|
|
seed=args.seed + j,
|
|
ref_imgs=ref_imgs_pil,
|
|
pe=args.pe,
|
|
siglip_inputs=siglip_inputs,
|
|
)
|
|
if args.concat_refs:
|
|
image_gen = horizontal_concat([image_gen, *ref_imgs])
|
|
|
|
if "save_dir" in data_dict:
|
|
config_save_path = os.path.join(args.save_path, data_dict["save_dir"] + f"_{j}.json")
|
|
image_save_path = os.path.join(args.save_path, data_dict["save_dir"] + f"_{j}.png")
|
|
else:
|
|
os.makedirs(args.save_path, exist_ok=True)
|
|
config_save_path = os.path.join(args.save_path, f"{i}_{j}.json")
|
|
image_save_path = os.path.join(args.save_path, f"{i}_{j}.png")
|
|
|
|
# save config and image
|
|
os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
|
|
image_gen.save(image_save_path)
|
|
# ensure the prompt and image_paths are saved in the config file
|
|
args.prompt = data_dict["prompt"]
|
|
args.image_paths = data_dict["image_paths"]
|
|
args_dict = vars(args)
|
|
with open(config_save_path, "w") as f:
|
|
json.dump(args_dict, f, indent=4)
|
|
|
|
if accelerator.is_main_process:
|
|
progress_bar.update(1)
|
|
if accelerator.is_main_process:
|
|
progress_bar.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = HfArgumentParser([InferenceArgs])
|
|
args = parser.parse_args_into_dataclasses()[0]
|
|
main(args)
|