mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-18 22:08:13 +00:00
85 lines
3.3 KiB
Python
85 lines
3.3 KiB
Python
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"
|
|
)
|