| from diffsynth import ModelManager, FluxImagePipeline, download_customized_models | |
| from modelscope import dataset_snapshot_download | |
| from examples.EntityControl.utils import visualize_masks | |
| from PIL import Image | |
| import torch | |
| def build_pipeline(): | |
| model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda", model_id_list=["FLUX.1-dev"]) | |
| model_manager.load_lora( | |
| download_customized_models( | |
| model_id="DiffSynth-Studio/Eligen", | |
| origin_file_path="model_bf16.safetensors", | |
| local_dir="models/lora/entity_control" | |
| ), | |
| lora_alpha=1 | |
| ) | |
| model_manager.load_lora( | |
| download_customized_models( | |
| model_id="iic/In-Context-LoRA", | |
| origin_file_path="visual-identity-design.safetensors", | |
| local_dir="models/lora/In-Context-LoRA" | |
| ), | |
| lora_alpha=1 | |
| ) | |
| pipe = FluxImagePipeline.from_model_manager(model_manager) | |
| return pipe | |
| def generate(pipe: FluxImagePipeline, source_image, target_image, mask, height, width, prompt, entity_prompt, image_save_path, mask_save_path, seed=0): | |
| input_mask = Image.new('RGB', (width * 2, height)) | |
| input_mask.paste(mask.resize((width, height), resample=Image.NEAREST).convert('RGB'), (width, 0)) | |
| input_image = Image.new('RGB', (width * 2, height)) | |
| input_image.paste(source_image.resize((width, height)).convert('RGB'), (0, 0)) | |
| input_image.paste(target_image.resize((width, height)).convert('RGB'), (width, 0)) | |
| image = pipe( | |
| prompt=prompt, | |
| input_image=input_image, | |
| cfg_scale=3.0, | |
| negative_prompt="", | |
| num_inference_steps=50, | |
| embedded_guidance=3.5, | |
| seed=seed, | |
| height=height, | |
| width=width * 2, | |
| eligen_entity_prompts=[entity_prompt], | |
| eligen_entity_masks=[input_mask], | |
| enable_eligen_on_negative=False, | |
| enable_eligen_inpaint=True, | |
| ) | |
| target_image = image.crop((width, 0, 2 * width, height)) | |
| target_image.save(image_save_path) | |
| visualize_masks(target_image, [mask], [entity_prompt], mask_save_path) | |
| return target_image | |
| pipe = build_pipeline() | |
| dataset_snapshot_download(dataset_id="DiffSynth-Studio/examples_in_diffsynth", local_dir="./", allow_file_pattern="data/examples/eligen/logo_transfer/*") | |
| prompt="The two-panel image showcases the joyful identity, with the left panel showing a rabbit graphic; [LEFT] while the right panel translates the design onto a shopping tote with the rabbit logo in black, held by a person in a market setting, emphasizing the brand's approachable and eco-friendly vibe." | |
| logo_prompt="a rabbit logo" | |
| logo_image = Image.open("data/examples/eligen/logo_transfer/source_image.png") | |
| target_image = Image.open("data/examples/eligen/logo_transfer/target_image.png") | |
| mask = Image.open("data/examples/eligen/logo_transfer/mask_1.png") | |
| generate( | |
| pipe, logo_image, target_image, mask, | |
| height=1024, width=1024, | |
| prompt=prompt, entity_prompt=logo_prompt, | |
| image_save_path="entity_transfer_1.png", | |
| mask_save_path="entity_transfer_with_mask_1.png" | |
| ) | |
| mask = Image.open("data/examples/eligen/logo_transfer/mask_2.png") | |
| generate( | |
| pipe, logo_image, target_image, mask, | |
| height=1024, width=1024, | |
| prompt=prompt, entity_prompt=logo_prompt, | |
| image_save_path="entity_transfer_2.png", | |
| mask_save_path="entity_transfer_with_mask_2.png" | |
| ) | |