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 [https://github.com/aihao2000/StableDiffusionReferenceOnly](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()