Spaces:
Sleeping
Sleeping
| import huggingface_hub | |
| import gradio as gr | |
| from stable_diffusion_reference_only.pipelines.stable_diffusion_reference_only_pipeline import ( | |
| StableDiffusionReferenceOnlyPipeline, | |
| ) | |
| import anime_segmentation | |
| from diffusers.schedulers import UniPCMultistepScheduler | |
| from PIL import Image | |
| import cv2 | |
| import numpy as np | |
| import os | |
| automatic_coloring_pipeline = StableDiffusionReferenceOnlyPipeline.from_pretrained( | |
| "AisingioroHao0/stable-diffusion-reference-only-automatic-coloring-0.1.2" | |
| ) | |
| automatic_coloring_pipeline.scheduler = UniPCMultistepScheduler.from_config( | |
| automatic_coloring_pipeline.scheduler.config | |
| ) | |
| segment_model = anime_segmentation.get_model( | |
| model_path=huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.ckpt") | |
| ) | |
| def character_segment(img): | |
| if img is None: | |
| return None | |
| img = anime_segmentation.character_segment(segment_model, img) | |
| img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) | |
| return img | |
| def color_inversion(img): | |
| if img is None: | |
| return None | |
| return 255 - img | |
| def get_line_art(img): | |
| if img is None: | |
| return None | |
| img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) | |
| img = cv2.adaptiveThreshold( | |
| img, | |
| 255, | |
| cv2.ADAPTIVE_THRESH_MEAN_C, | |
| cv2.THRESH_BINARY, | |
| blockSize=5, | |
| C=7, | |
| ) | |
| img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) | |
| return img | |
| def inference(prompt, blueprint, num_inference_steps): | |
| if prompt is None or blueprint is None: | |
| return None | |
| return np.array( | |
| automatic_coloring_pipeline( | |
| prompt=Image.fromarray(prompt), | |
| blueprint=Image.fromarray(blueprint), | |
| num_inference_steps=num_inference_steps, | |
| ).images[0] | |
| ) | |
| def automatic_coloring(prompt, blueprint, num_inference_steps): | |
| if prompt is None or blueprint is None: | |
| return None | |
| blueprint = color_inversion(blueprint) | |
| return inference(prompt, blueprint, num_inference_steps) | |
| def style_transfer(prompt, blueprint, num_inference_steps): | |
| if prompt is None or blueprint is None: | |
| return None | |
| prompt = character_segment(prompt) | |
| blueprint = character_segment(blueprint) | |
| blueprint = get_line_art(blueprint) | |
| blueprint = color_inversion(blueprint) | |
| return inference(prompt, blueprint, num_inference_steps) | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # Stable Diffusion Reference Only Automatic Coloring 0.1.2\n\n | |
| demo for [<svg height="32" aria-hidden="true" viewBox="0 0 16 16" version="1.1" width="32" data-view-component="true" class="octicon octicon-mark-github v-align-middle color-fg-default"> | |
| <path d="M8 0c4.42 0 8 3.58 8 8a8.013 8.013 0 0 1-5.45 7.59c-.4.08-.55-.17-.55-.38 0-.27.01-1.13.01-2.2 0-.75-.25-1.23-.54-1.48 1.78-.2 3.65-.88 3.65-3.95 0-.88-.31-1.59-.82-2.15.08-.2.36-1.02-.08-2.12 0 0-.67-.22-2.2.82-.64-.18-1.32-.27-2-.27-.68 0-1.36.09-2 .27-1.53-1.03-2.2-.82-2.2-.82-.44 1.1-.16 1.92-.08 2.12-.51.56-.82 1.28-.82 2.15 0 3.06 1.86 3.75 3.64 3.95-.23.2-.44.55-.51 1.07-.46.21-1.61.55-2.33-.66-.15-.24-.6-.83-1.23-.82-.67.01-.27.38.01.53.34.19.73.9.82 1.13.16.45.68 1.31 2.69.94 0 .67.01 1.3.01 1.49 0 .21-.15.45-.55.38A7.995 7.995 0 0 1 0 8c0-4.42 3.58-8 8-8Z"></path> | |
| </svg>](https://github.com/aihao2000/StableDiffusionReferenceOnly) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt_input_compoent = gr.Image(shape=(512, 512), label="prompt") | |
| prompt_character_segment_button = gr.Button( | |
| "character segment", | |
| ) | |
| prompt_character_segment_button.click( | |
| character_segment, | |
| inputs=prompt_input_compoent, | |
| outputs=prompt_input_compoent, | |
| ) | |
| with gr.Column(): | |
| blueprint_input_compoent = gr.Image(shape=(512, 512), label="blueprint") | |
| blueprint_character_segment_button = gr.Button("character segment") | |
| blueprint_character_segment_button.click( | |
| character_segment, | |
| inputs=blueprint_input_compoent, | |
| outputs=blueprint_input_compoent, | |
| ) | |
| get_line_art_button = gr.Button( | |
| "get line art", | |
| ) | |
| get_line_art_button.click( | |
| get_line_art, | |
| inputs=blueprint_input_compoent, | |
| outputs=blueprint_input_compoent, | |
| ) | |
| color_inversion_button = gr.Button( | |
| "color inversion", | |
| ) | |
| color_inversion_button.click( | |
| color_inversion, | |
| inputs=blueprint_input_compoent, | |
| outputs=blueprint_input_compoent, | |
| ) | |
| with gr.Column(): | |
| result_output_component = gr.Image(shape=(512, 512), label="result") | |
| num_inference_steps_input_component = gr.Number( | |
| 20, label="num inference steps", minimum=1, maximum=1000, step=1 | |
| ) | |
| inference_button = gr.Button("inference") | |
| inference_button.click( | |
| inference, | |
| inputs=[ | |
| prompt_input_compoent, | |
| blueprint_input_compoent, | |
| num_inference_steps_input_component, | |
| ], | |
| outputs=result_output_component, | |
| ) | |
| automatic_coloring_button = gr.Button("automatic coloring") | |
| automatic_coloring_button.click( | |
| automatic_coloring, | |
| inputs=[ | |
| prompt_input_compoent, | |
| blueprint_input_compoent, | |
| num_inference_steps_input_component, | |
| ], | |
| outputs=result_output_component, | |
| ) | |
| style_transfer_button = gr.Button("style transfer") | |
| style_transfer_button.click( | |
| style_transfer, | |
| inputs=[ | |
| prompt_input_compoent, | |
| blueprint_input_compoent, | |
| num_inference_steps_input_component, | |
| ], | |
| outputs=result_output_component, | |
| ) | |
| with gr.Row(): | |
| gr.Examples( | |
| examples=[ | |
| [ | |
| os.path.join( | |
| os.path.dirname(__file__), "README.assets", "3x9_prompt.png" | |
| ), | |
| os.path.join( | |
| os.path.dirname(__file__), "README.assets", "3x9_blueprint.png" | |
| ), | |
| ], | |
| ], | |
| inputs=[prompt_input_compoent, blueprint_input_compoent], | |
| outputs=result_output_component, | |
| fn=lambda x, y: None, | |
| cache_examples=True, | |
| ) | |
| demo.launch() | |