|
import gradio as gr |
|
import torch |
|
import numpy as np |
|
import spaces |
|
from utils import utils, tools, preprocess |
|
|
|
BASE_MODEL_REPO_ID = "neta-art/neta-xl-2.0" |
|
BASE_MODEL_FILENAME = "neta-xl-v2.fp16.safetensors" |
|
VAE_PATH = "madebyollin/sdxl-vae-fp16-fix" |
|
CONTROLNEXT_REPO_ID = "Eugeoter/controlnext-sdxl-anime-canny" |
|
CACHE_DIR = None |
|
|
|
DEFAULT_PROMPT = "" |
|
DEFAULT_NEGATIVE_PROMPT = "worst quality, abstract, clumsy pose, deformed hand, dynamic malformation, fused fingers, extra digits, fewer digits, fewer fingers, extra fingers, extra arm, missing arm, extra leg, missing leg, signature, artist name, multi views, disfigured, ugly" |
|
|
|
|
|
def ui(): |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
pipeline = tools.get_pipeline( |
|
pretrained_model_name_or_path=BASE_MODEL_REPO_ID, |
|
unet_model_name_or_path=CONTROLNEXT_REPO_ID, |
|
controlnet_model_name_or_path=CONTROLNEXT_REPO_ID, |
|
vae_model_name_or_path=VAE_PATH, |
|
load_weight_increasement=True, |
|
device=device, |
|
hf_cache_dir=CACHE_DIR, |
|
use_safetensors=True, |
|
) |
|
|
|
schedulers = ['Euler A', 'UniPC', 'Euler', 'DDIM', 'DDPM'] |
|
|
|
css = """ |
|
#col-container { |
|
margin: 0 auto; |
|
max-width: 520px; |
|
} |
|
""" |
|
|
|
with gr.Blocks(css=css) as demo: |
|
gr.Markdown(f""" |
|
# [ControlNeXt-SDXL](https://github.com/dvlab-research/ControlNeXt) Demo (Anime Canny) |
|
Base model: [Neta-Art-XL-2.0](https://civitai.com/models/410737/neta-art-xl) |
|
""") |
|
with gr.Row(): |
|
with gr.Column(scale=9): |
|
prompt = gr.Textbox(label='Prompt', value=DEFAULT_PROMPT, lines=3, placeholder='prompt', container=False) |
|
negative_prompt = gr.Textbox(label='Negative Prompt', value=DEFAULT_NEGATIVE_PROMPT, lines=3, placeholder='negative prompt', container=False) |
|
with gr.Column(scale=1): |
|
generate_button = gr.Button("Generate", variant='primary', min_width=96) |
|
with gr.Row(): |
|
with gr.Column(scale=1): |
|
with gr.Row(): |
|
control_image = gr.Image( |
|
value=None, |
|
label='Condition', |
|
sources=['upload'], |
|
type='pil', |
|
height=512, |
|
image_mode='RGB', |
|
format='png', |
|
show_download_button=True, |
|
show_share_button=True, |
|
) |
|
with gr.Accordion(label='Preprocess', open=True): |
|
with gr.Row(): |
|
threshold1 = gr.Slider(minimum=-1, maximum=255, step=1, value=-1, label='Threshold 1', info='-1 for auto') |
|
threshold2 = gr.Slider(minimum=-1, maximum=255, step=1, value=-1, label='Threshold 2', info='-1 for auto') |
|
process_button = gr.Button("Process", variant='primary', min_width=96, scale=0) |
|
with gr.Row(): |
|
scheduler = gr.Dropdown( |
|
label='Scheduler', |
|
choices=schedulers, |
|
value='Euler A', |
|
multiselect=False, |
|
allow_custom_value=False, |
|
filterable=True, |
|
) |
|
num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, value=28, label='Steps') |
|
with gr.Row(): |
|
cfg_scale = gr.Slider(minimum=1, maximum=30, step=1, value=7.5, label='CFG Scale') |
|
controlnet_scale = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label='ControlNet Scale') |
|
with gr.Row(): |
|
seed = gr.Number(label='Seed', step=1, precision=0, value=-1) |
|
with gr.Column(scale=1): |
|
with gr.Row(): |
|
output = gr.Gallery( |
|
label='Output', |
|
value=None, |
|
object_fit='scale-down', |
|
columns=4, |
|
height=512, |
|
show_download_button=True, |
|
show_share_button=True, |
|
) |
|
|
|
with gr.Row(): |
|
examples = gr.Examples( |
|
label='Examples', |
|
examples=[ |
|
[ |
|
'best quality, 1girl, solo, open hand, outdoors, indoor, cute, young, cat, cat ear, glasses', |
|
'examples/example_1.jpg', |
|
], |
|
[ |
|
'best quality, 1 komeiji koishi, solo, the pose, indoors, smile', |
|
'examples/example_2.jpg', |
|
] |
|
], |
|
inputs=[ |
|
prompt, |
|
control_image, |
|
], |
|
cache_examples=False, |
|
) |
|
|
|
@spaces.GPU |
|
def generate( |
|
prompt, |
|
control_image, |
|
negative_prompt, |
|
cfg_scale, |
|
controlnet_scale, |
|
num_inference_steps, |
|
scheduler, |
|
seed, |
|
): |
|
pipeline.scheduler = tools.get_scheduler(scheduler, pipeline.scheduler.config) |
|
generator = torch.Generator(device=device).manual_seed(max(0, min(seed, np.iinfo(np.int32).max))) if seed != -1 else None |
|
|
|
if control_image is None: |
|
raise gr.Error('Please upload an image.') |
|
width, height = utils.around_reso(control_image.width, control_image.height, reso=1024, max_width=2048, max_height=2048, divisible=32) |
|
control_image = control_image.resize((width, height)).convert('RGB') |
|
|
|
with torch.autocast(device): |
|
output_images = pipeline.__call__( |
|
prompt=prompt, |
|
negative_prompt=negative_prompt, |
|
controlnet_image=control_image, |
|
controlnet_scale=controlnet_scale, |
|
width=width, |
|
height=height, |
|
generator=generator, |
|
guidance_scale=cfg_scale, |
|
num_inference_steps=num_inference_steps, |
|
).images |
|
|
|
return output_images |
|
|
|
def process( |
|
image, |
|
threshold1, |
|
threshold2, |
|
): |
|
threshold1 = None if threshold1 == -1 else threshold1 |
|
threshold2 = None if threshold2 == -1 else threshold2 |
|
return preprocess.canny_extractor(image, threshold1, threshold2) |
|
|
|
generate_button.click( |
|
fn=generate, |
|
inputs=[prompt, control_image, negative_prompt, cfg_scale, controlnet_scale, num_inference_steps, scheduler, seed], |
|
outputs=[output], |
|
) |
|
|
|
process_button.click( |
|
fn=process, |
|
inputs=[control_image, threshold1, threshold2], |
|
outputs=[control_image], |
|
) |
|
|
|
return demo |
|
|
|
|
|
if __name__ == '__main__': |
|
demo = ui() |
|
demo.queue().launch() |
|
|
