'''
sudo apt-get update && sudo apt-get install git-lfs cbm ffmpeg
huggingface-cli login
'''
import sys
sys.path.append('../')
import spaces
import torch
import random
import numpy as np
from PIL import Image
import gradio as gr
from huggingface_hub import hf_hub_download
from transformers import AutoModelForImageSegmentation
from torchvision import transforms
from pipeline import InstantCharacterFluxPipeline
# global variable
MAX_SEED = np.iinfo(np.int32).max
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if str(device).__contains__("cuda") else torch.float32
# pre-trained weights
ip_adapter_path = hf_hub_download(repo_id="tencent/InstantCharacter", filename="instantcharacter_ip-adapter.bin")
base_model = 'black-forest-labs/FLUX.1-dev'
image_encoder_path = 'google/siglip-so400m-patch14-384'
image_encoder_2_path = 'facebook/dinov2-giant'
birefnet_path = 'ZhengPeng7/BiRefNet'
makoto_style_lora_path = hf_hub_download(repo_id="InstantX/FLUX.1-dev-LoRA-Makoto-Shinkai", filename="Makoto_Shinkai_style.safetensors")
ghibli_style_lora_path = hf_hub_download(repo_id="InstantX/FLUX.1-dev-LoRA-Ghibli", filename="ghibli_style.safetensors")
# init InstantCharacter pipeline
pipe = InstantCharacterFluxPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
pipe.to(device)
# load InstantCharacter
pipe.init_adapter(
image_encoder_path=image_encoder_path,
image_encoder_2_path=image_encoder_2_path,
subject_ipadapter_cfg=dict(subject_ip_adapter_path=ip_adapter_path, nb_token=1024),
)
#pipe.enable_sequential_cpu_offload()
# load matting model
birefnet = AutoModelForImageSegmentation.from_pretrained(birefnet_path, trust_remote_code=True)
birefnet.to('cuda')
birefnet.eval()
birefnet_transform_image = transforms.Compose([
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def remove_bkg(subject_image):
def infer_matting(img_pil):
input_images = birefnet_transform_image(img_pil).unsqueeze(0).to('cuda')
with torch.no_grad():
preds = birefnet(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(img_pil.size)
mask = np.array(mask)
mask = mask[..., None]
return mask
def get_bbox_from_mask(mask, th=128):
height, width = mask.shape[:2]
x1, y1, x2, y2 = 0, 0, width - 1, height - 1
sample = np.max(mask, axis=0)
for idx in range(width):
if sample[idx] >= th:
x1 = idx
break
sample = np.max(mask[:, ::-1], axis=0)
for idx in range(width):
if sample[idx] >= th:
x2 = width - 1 - idx
break
sample = np.max(mask, axis=1)
for idx in range(height):
if sample[idx] >= th:
y1 = idx
break
sample = np.max(mask[::-1], axis=1)
for idx in range(height):
if sample[idx] >= th:
y2 = height - 1 - idx
break
x1 = np.clip(x1, 0, width-1).round().astype(np.int32)
y1 = np.clip(y1, 0, height-1).round().astype(np.int32)
x2 = np.clip(x2, 0, width-1).round().astype(np.int32)
y2 = np.clip(y2, 0, height-1).round().astype(np.int32)
return [x1, y1, x2, y2]
def pad_to_square(image, pad_value = 255, random = False):
'''
image: np.array [h, w, 3]
'''
H,W = image.shape[0], image.shape[1]
if H == W:
return image
padd = abs(H - W)
if random:
padd_1 = int(np.random.randint(0,padd))
else:
padd_1 = int(padd / 2)
padd_2 = padd - padd_1
if H > W:
pad_param = ((0,0),(padd_1,padd_2),(0,0))
else:
pad_param = ((padd_1,padd_2),(0,0),(0,0))
image = np.pad(image, pad_param, 'constant', constant_values=pad_value)
return image
salient_object_mask = infer_matting(subject_image)[..., 0]
x1, y1, x2, y2 = get_bbox_from_mask(salient_object_mask)
subject_image = np.array(subject_image)
salient_object_mask[salient_object_mask > 128] = 255
salient_object_mask[salient_object_mask < 128] = 0
sample_mask = np.concatenate([salient_object_mask[..., None]]*3, axis=2)
obj_image = sample_mask / 255 * subject_image + (1 - sample_mask / 255) * 255
crop_obj_image = obj_image[y1:y2, x1:x2]
crop_pad_obj_image = pad_to_square(crop_obj_image, 255)
subject_image = Image.fromarray(crop_pad_obj_image.astype(np.uint8))
return subject_image
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
if randomize_seed:
seed = random.randint(0, MAX_SEED)
return seed
def get_example():
case = [
[
"./assets/girl.jpg",
"A girl is playing a guitar in street",
0.9,
'Makoto Shinkai style',
],
[
"./assets/boy.jpg",
"A boy is riding a bike in snow",
0.9,
'Makoto Shinkai style',
],
]
return case
def run_for_examples(source_image, prompt, scale, style_mode):
return create_image(
input_image=source_image,
prompt=prompt,
scale=scale,
guidance_scale=3.5,
num_inference_steps=28,
seed=123456,
style_mode=style_mode,
)
@spaces.GPU
def create_image(input_image,
prompt,
scale,
guidance_scale,
num_inference_steps,
seed,
style_mode=None):
input_image = remove_bkg(input_image)
if style_mode is None:
images = pipe(
prompt=prompt,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
width=1024,
height=1024,
subject_image=input_image,
subject_scale=scale,
generator=torch.manual_seed(seed),
).images
else:
if style_mode == 'Makoto Shinkai style':
lora_file_path = makoto_style_lora_path
trigger = 'Makoto Shinkai style'
elif style_mode == 'Ghibli style':
lora_file_path = ghibli_style_lora_path
trigger = 'ghibli style'
images = pipe.with_style_lora(
lora_file_path=lora_file_path,
trigger=trigger,
prompt=prompt,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
width=1024,
height=1024,
subject_image=input_image,
subject_scale=scale,
generator=torch.manual_seed(seed),
).images
return images
# Description
title = r"""
<h1 align="center">InstantCharacter : Personalize Any Characters with a Scalable Diffusion Transformer Framework</h1>
"""
description = r"""
<b>Official 🤗 Gradio demo</b> for <a href='https://instantcharacter.github.io/' target='_blank'><b>InstantCharacter : Personalize Any Characters with a Scalable Diffusion Transformer Framework</b></a>.<br>
How to use:<br>
1. Upload a character image, removing background would be preferred.
2. Enter a text prompt to describe what you hope the chracter does.
3. Click the <b>Submit</b> button to begin customization.
4. Share your custimized photo with your friends and enjoy! 😊
"""
article = r"""
---
📝 **Citation**
<br>
If our work is helpful for your research or applications, please cite us via:
```bibtex
TBD
```
📧 **Contact**
<br>
If you have any questions, please feel free to open an issue.
"""
block = gr.Blocks(css="footer {visibility: hidden}").queue(max_size=10, api_open=False)
with block:
# description
gr.Markdown(title)
gr.Markdown(description)
with gr.Tabs():
with gr.Row():
with gr.Column():
with gr.Row():
with gr.Column():
image_pil = gr.Image(label="Source Image", type='pil')
prompt = gr.Textbox(label="Prompt", value="a character is riding a bike in snow")
scale = gr.Slider(minimum=0, maximum=1.5, step=0.01,value=1.0, label="Scale")
style_mode = gr.Dropdown(label='Style', choices=[None, 'Makoto Shinkai style', 'Ghibli style'], value='Makoto Shinkai style')
with gr.Accordion(open=False, label="Advanced Options"):
guidance_scale = gr.Slider(minimum=1,maximum=7.0, step=0.01,value=3.5, label="guidance scale")
num_inference_steps = gr.Slider(minimum=5,maximum=50.0, step=1.0,value=28, label="num inference steps")
seed = gr.Slider(minimum=-1000000, maximum=1000000, value=123456, step=1, label="Seed Value")
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
generate_button = gr.Button("Generate Image")
with gr.Column():
generated_image = gr.Gallery(label="Generated Image")
generate_button.click(
fn=randomize_seed_fn,
inputs=[seed, randomize_seed],
outputs=seed,
queue=False,
api_name=False,
).then(
fn=create_image,
inputs=[image_pil,
prompt,
scale,
guidance_scale,
num_inference_steps,
seed,
style_mode,
],
outputs=[generated_image])
gr.Examples(
examples=get_example(),
inputs=[image_pil, prompt, scale, style_mode],
fn=run_for_examples,
outputs=[generated_image],
cache_examples=True,
)
gr.Markdown(article)
block.launch(share = True)