mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-18 22:08:13 +00:00
55 lines
2.3 KiB
Python
55 lines
2.3 KiB
Python
import torch
|
|
from diffsynth import ModelManager, FluxImagePipeline, download_customized_models, FluxImageLoraPipeline
|
|
from examples.EntityControl.utils import visualize_masks
|
|
import os
|
|
import json
|
|
from PIL import Image
|
|
|
|
# lora_path = download_customized_models(
|
|
# model_id="DiffSynth-Studio/ArtAug-lora-FLUX.1dev-v1",
|
|
# origin_file_path="merged_lora.safetensors",
|
|
# local_dir="models/lora"
|
|
# )[0]
|
|
|
|
lora_path = '/root/model_bf16.safetensors'
|
|
model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda")
|
|
model_manager.load_models([
|
|
"t2i_models/FLUX/FLUX.1-dev/text_encoder/model.safetensors",
|
|
"t2i_models/FLUX/FLUX.1-dev/text_encoder_2",
|
|
"t2i_models/FLUX/FLUX.1-dev/ae.safetensors",
|
|
"t2i_models/FLUX/FLUX.1-dev/flux1-dev.safetensors"
|
|
])
|
|
model_manager.load_lora(lora_path, lora_alpha=1.)
|
|
|
|
pipe = FluxImagePipeline.from_model_manager(model_manager)
|
|
|
|
mask_dir = '/mnt/nas1/zhanghong/DiffSynth-Studio/workdirs/tmp_mask'
|
|
image_shape = 1024
|
|
guidance = 3.5
|
|
cfg = 3.0
|
|
negative_prompt = "worst quality, low quality, monochrome, zombie, interlocked fingers, Aissist, cleavage, nsfw,"
|
|
names = ['row_2_1']
|
|
seeds = [0]
|
|
# use this to apply regional attention in negative prompt prediction for better results with more time
|
|
use_seperated_negtive_prompt = False
|
|
for name, seed in zip(names, seeds):
|
|
out_dir = f'workdirs/entity_control/{name}'
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
cur_dir = os.path.join(mask_dir, name)
|
|
metas = json.load(open(os.path.join(mask_dir, name, 'prompts.json')))
|
|
for seed in range(3, 10):
|
|
prompt = metas['global_prompt']
|
|
mask_prompts = metas['mask_prompts']
|
|
masks = [Image.open(os.path.join(mask_dir, name, f"{mask_idx}.png")).resize((image_shape, image_shape), resample=Image.NEAREST) for mask_idx in range(len(mask_prompts))]
|
|
torch.manual_seed(seed)
|
|
image = pipe(
|
|
prompt=prompt,
|
|
cfg_scale=cfg,
|
|
negative_prompt=negative_prompt,
|
|
num_inference_steps=50, embedded_guidance=guidance, height=image_shape, width=image_shape,
|
|
entity_prompts=mask_prompts, entity_masks=masks,
|
|
use_seperated_negtive_prompt=use_seperated_negtive_prompt
|
|
)
|
|
use_sep = f'_sepneg' if use_seperated_negtive_prompt else ''
|
|
visualize_masks(image, masks, mask_prompts, os.path.join(out_dir, f"{name}_{seed}{use_sep}.png"))
|