|
import spaces |
|
import time |
|
import os |
|
|
|
import gradio as gr |
|
import torch |
|
from einops import rearrange |
|
from PIL import Image |
|
|
|
from flux.cli import SamplingOptions |
|
from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack |
|
from flux.util import load_ae, load_clip, load_flow_model, load_t5 |
|
from pulid.pipeline_flux import PuLIDPipeline |
|
from pulid.utils import resize_numpy_image_long |
|
|
|
|
|
def get_models(name: str, device: torch.device, offload: bool): |
|
t5 = load_t5(device, max_length=128) |
|
clip = load_clip(device) |
|
model = load_flow_model(name, device="cpu" if offload else device) |
|
model.eval() |
|
ae = load_ae(name, device="cpu" if offload else device) |
|
return model, ae, t5, clip |
|
|
|
|
|
class FluxGenerator: |
|
def __init__(self): |
|
self.device = torch.device('cuda') |
|
self.offload = False |
|
self.model_name = 'flux-dev' |
|
self.model, self.ae, self.t5, self.clip = get_models( |
|
self.model_name, |
|
device=self.device, |
|
offload=self.offload, |
|
) |
|
self.pulid_model = PuLIDPipeline(self.model, 'cuda', weight_dtype=torch.bfloat16) |
|
self.pulid_model.load_pretrain() |
|
|
|
|
|
flux_generator = FluxGenerator() |
|
|
|
|
|
@spaces.GPU |
|
@torch.inference_mode() |
|
def generate_image( |
|
width, |
|
height, |
|
num_steps, |
|
start_step, |
|
guidance, |
|
seed, |
|
prompt, |
|
id_image=None, |
|
id_weight=1.0, |
|
neg_prompt="", |
|
true_cfg=1.0, |
|
timestep_to_start_cfg=1, |
|
max_sequence_length=128, |
|
): |
|
flux_generator.t5.max_length = max_sequence_length |
|
|
|
seed = int(seed) |
|
if seed == -1: |
|
seed = None |
|
|
|
opts = SamplingOptions( |
|
prompt=prompt, |
|
width=width, |
|
height=height, |
|
num_steps=num_steps, |
|
guidance=guidance, |
|
seed=seed, |
|
) |
|
|
|
if opts.seed is None: |
|
opts.seed = torch.Generator(device="cpu").seed() |
|
|
|
t0 = time.perf_counter() |
|
|
|
use_true_cfg = abs(true_cfg - 1.0) > 1e-2 |
|
|
|
if id_image is not None: |
|
id_image = resize_numpy_image_long(id_image, 1024) |
|
id_embeddings, uncond_id_embeddings = flux_generator.pulid_model.get_id_embedding(id_image, cal_uncond=use_true_cfg) |
|
else: |
|
id_embeddings = None |
|
uncond_id_embeddings = None |
|
|
|
|
|
x = get_noise( |
|
1, |
|
opts.height, |
|
opts.width, |
|
device=flux_generator.device, |
|
dtype=torch.bfloat16, |
|
seed=opts.seed, |
|
) |
|
timesteps = get_schedule( |
|
opts.num_steps, |
|
x.shape[-1] * x.shape[-2] // 4, |
|
shift=True, |
|
) |
|
|
|
if flux_generator.offload: |
|
flux_generator.t5, flux_generator.clip = flux_generator.t5.to(flux_generator.device), flux_generator.clip.to(flux_generator.device) |
|
inp = prepare(t5=flux_generator.t5, clip=flux_generator.clip, img=x, prompt=opts.prompt) |
|
inp_neg = prepare(t5=flux_generator.t5, clip=flux_generator.clip, img=x, prompt=neg_prompt) if use_true_cfg else None |
|
|
|
|
|
if flux_generator.offload: |
|
flux_generator.t5, flux_generator.clip = flux_generator.t5.cpu(), flux_generator.clip.cpu() |
|
torch.cuda.empty_cache() |
|
flux_generator.model = flux_generator.model.to(flux_generator.device) |
|
|
|
|
|
x = denoise( |
|
flux_generator.model, **inp, timesteps=timesteps, guidance=opts.guidance, id=id_embeddings, id_weight=id_weight, |
|
start_step=start_step, uncond_id=uncond_id_embeddings, true_cfg=true_cfg, |
|
timestep_to_start_cfg=timestep_to_start_cfg, |
|
neg_txt=inp_neg["txt"] if use_true_cfg else None, |
|
neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None, |
|
neg_vec=inp_neg["vec"] if use_true_cfg else None, |
|
) |
|
|
|
|
|
if flux_generator.offload: |
|
flux_generator.model.cpu() |
|
torch.cuda.empty_cache() |
|
flux_generator.ae.decoder.to(x.device) |
|
|
|
|
|
x = unpack(x.float(), opts.height, opts.width) |
|
with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16): |
|
x = flux_generator.ae.decode(x) |
|
|
|
if flux_generator.offload: |
|
flux_generator.ae.decoder.cpu() |
|
torch.cuda.empty_cache() |
|
|
|
t1 = time.perf_counter() |
|
|
|
|
|
x = x.clamp(-1, 1) |
|
x = rearrange(x[0], "c h w -> h w c") |
|
|
|
img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy()) |
|
return img, str(opts.seed), flux_generator.pulid_model.debug_img_list |
|
|
|
|
|
css = """ |
|
footer { |
|
visibility: hidden; |
|
} |
|
""" |
|
|
|
|
|
def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu", |
|
offload: bool = False): |
|
|
|
with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo: |
|
gr.Markdown("### 'AI ํฌํ ์ง๋'์ด์ฉ ์๋ด: 1) '์คํ์ผ'์ค ํ๋๋ฅผ ์ ํ. 2) ์น์บ ์ ํด๋ฆญํ๊ณ ์ผ๊ตด์ด ๋ณด์ด๋ฉด ์นด๋ฉ๋ผ ๋ฒํผ ํด๋ฆญ. 3) '์์ฑ' ๋ฒํผ์ ํด๋ฆญํ๊ณ ๊ธฐ๋ค๋ฆฌ๋ฉด ๋ฉ๋๋ค.") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
prompt = gr.Textbox(label="ํ๋กฌํํธ", value="์ด์ํ, ์๊ฐ, ์ํ์ ") |
|
id_image = gr.Image(label="ID ์ด๋ฏธ์ง", sources=["webcam", "upload"], type="numpy") |
|
generate_btn = gr.Button("์์ฑ") |
|
|
|
with gr.Column(): |
|
output_image = gr.Image(label="์์ฑ๋ ์ด๋ฏธ์ง") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
gr.Markdown("### ์คํ์ผ") |
|
|
|
all_examples = [ |
|
["์ฐ์ฃผ ์ฌํI", 'I am an astronaut on a spacewalk. There is no helmet, and my face is visible. The background is Earth & starship as seen from space shuttle.', 'example_inputs/1.webp'], |
|
["์ฐ์ฃผ ์ฌํII", 'I am an astronaut on a spacewalk. There is no helmet, and my face is visible. The background is Earth & starship as seen from space shuttle.I am holding sign with glowing green text "I Love Mom"', 'example_inputs/2.webp'], |
|
["์ฐ์ฃผ์ ์กฐ์ข
", 'I am wearing a spacesuit and have become an astronaut. I am piloting a spacecraft. Through the spacecraft's window, I can see outer space.', 'example_inputs/7.webp'], |
|
["ํ์ฑ ํํ", 'I am wearing a spacesuit and have become an astronaut walking on Mars. I'm not wearing a helmet. I'm looking straight ahead. The background is a desolate area of Mars, and a space rover and a space station can be seen.', 'example_inputs/5.webp'], |
|
["์์ด์ธ๋งจ ๋ณ์ ", 'I am an "IRON MAN"', 'example_inputs/4.webp'], |
|
["์คํ์ด๋๋งจ", 'I am an "spider MAN"', 'example_inputs/6.webp'], |
|
["๋งํ ์ฃผ์ธ๊ณต", 'portrait, pixar style', 'example_inputs/8.webp'], |
|
["์๋์ฐ๋จผ", 'I am an "wonder woman"', 'example_inputs/9.webp'], |
|
["์นด์ฐ๋ณด์ด", 'Cowboy, american comics style', 'example_inputs/10.webp'], |
|
["๋ด๊ฐ ์ด๋ฅธ์ด ๋๋ฉด", 'profile photo of a 40-year-old Adult Looking straight ahead, wear suite', 'example_inputs/3.webp'], |
|
] |
|
|
|
example_gallery = gr.Gallery( |
|
[example[2] for example in all_examples], |
|
label="์คํ์ผ ์์", |
|
elem_id="gallery", |
|
columns=5, |
|
rows=2, |
|
object_fit="contain", |
|
height="auto" |
|
) |
|
|
|
def fill_example(evt: gr.SelectData): |
|
return all_examples[evt.index][1] |
|
|
|
example_gallery.select( |
|
fill_example, |
|
None, |
|
[prompt], |
|
) |
|
|
|
generate_btn.click( |
|
fn=lambda *args: generate_image(*args)[0], # ์ฒซ ๋ฒ์งธ ํญ๋ชฉ(์ด๋ฏธ์ง)๋ง ๋ฐํ |
|
inputs=[ |
|
gr.Slider(256, 1536, 896, step=16, visible=False), # width |
|
gr.Slider(256, 1536, 1152, step=16, visible=False), # height |
|
gr.Slider(1, 20, 20, step=1, visible=False), # num_steps |
|
gr.Slider(0, 10, 0, step=1, visible=False), # start_step |
|
gr.Slider(1.0, 10.0, 4, step=0.1, visible=False), # guidance |
|
gr.Textbox(-1, visible=False), # seed |
|
prompt, |
|
id_image, |
|
gr.Slider(0.0, 3.0, 1, step=0.05, visible=False), # id_weight |
|
gr.Textbox("Low quality, worst quality, text, signature, watermark, extra limbs", visible=False), # neg_prompt |
|
gr.Slider(1.0, 10.0, 1, step=0.1, visible=False), # true_cfg |
|
gr.Slider(0, 20, 1, step=1, visible=False), # timestep_to_start_cfg |
|
gr.Slider(128, 512, 128, step=128, visible=False), # max_sequence_length |
|
], |
|
outputs=[output_image], |
|
) |
|
|
|
return demo |
|
|
|
|
|
if __name__ == "__main__": |
|
import argparse |
|
|
|
parser = argparse.ArgumentParser(description="PuLID for FLUX.1-dev") |
|
parser.add_argument("--name", type=str, default="flux-dev", choices=list('flux-dev'), |
|
help="ํ์ฌ๋ flux-dev๋ง ์ง์ํฉ๋๋ค") |
|
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", |
|
help="์ฌ์ฉํ ๋๋ฐ์ด์ค") |
|
parser.add_argument("--offload", action="store_true", help="์ฌ์ฉํ์ง ์์ ๋ ๋ชจ๋ธ์ CPU๋ก ์ฎ๊น๋๋ค") |
|
parser.add_argument("--port", type=int, default=8080, help="์ฌ์ฉํ ํฌํธ") |
|
parser.add_argument("--dev", action='store_true', help="๊ฐ๋ฐ ๋ชจ๋") |
|
parser.add_argument("--pretrained_model", type=str, help='๊ฐ๋ฐ์ฉ') |
|
args = parser.parse_args() |
|
|
|
import huggingface_hub |
|
huggingface_hub.login(os.getenv('HF_TOKEN')) |
|
|
|
demo = create_demo(args, args.name, args.device, args.offload) |
|
demo.launch(share=True) |
|
|