add hair color app

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
 
 from app_base import create_demo as create_demo_face
+from app_haircolor import create_demo as create_demo_haircolor
 
 with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():
         with gr.Tab(label="Face"):
             create_demo_face()
+        with gr.Tab(label="Hair Color"):
+            create_demo_haircolor()
 
 demo.launch()

app_base.py

@@ -24,14 +24,6 @@ DEFAULT_CATEGORY = "face"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-os.system("pip freeze")
-if not os.path.exists('GFPGANv1.4.pth'):
-    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
-if not os.path.exists('realesr-general-x4v3.pth'):
-    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
-
-os.makedirs('output', exist_ok=True)
-
 def create_demo() -> gr.Blocks:
     model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
     model_path = 'realesr-general-x4v3.pth'
@@ -40,7 +32,7 @@ def create_demo() -> gr.Blocks:
 
     face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=1, arch='clean', channel_multiplier=2)
 
-    @spaces.GPU(duration=15)
+    @spaces.GPU(duration=10)
     def image_to_image(
         input_image: Image,
         input_image_prompt: str,

app_haircolor.py

@@ -0,0 +1,142 @@
+import spaces
+import gradio as gr
+import time
+import torch
+import os
+import numpy as np
+import cv2
+
+from PIL import Image
+from inversion_run_base import run as base_run
+from segment_utils import(
+    segment_image,
+    restore_result,
+)
+from gfpgan.utils import GFPGANer
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from realesrgan.utils import RealESRGANer
+
+
+DEFAULT_SRC_PROMPT = "a woman"
+DEFAULT_EDIT_PROMPT = "a woman, with blue hair, 8k, high quality"
+
+DEFAULT_CATEGORY = "hair"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def create_demo() -> gr.Blocks:
+    model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+    model_path = 'realesr-general-x4v3.pth'
+    half = True if torch.cuda.is_available() else False
+    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+    @spaces.GPU(duration=10)
+    def image_to_image(
+        input_image: Image,
+        input_image_prompt: str,
+        edit_prompt: str,
+        seed: int,
+        w1: float,
+        num_steps: int,
+        start_step: int,
+        guidance_scale: float,
+        generate_size: int,
+        adapter_weights: float,
+    ):
+        w2 = 1.0
+        run_task_time = 0
+        time_cost_str = ''
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        run_model = base_run
+        res_image = run_model(
+            input_image,
+            input_image_prompt,
+            edit_prompt,
+            generate_size,
+            seed,
+            w1,
+            w2,
+            num_steps,
+            start_step,
+            guidance_scale,
+            adapter_weights,
+        )
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        enhanced_image = enhance(res_image)
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+
+        return enhanced_image, res_image, time_cost_str
+
+    def get_time_cost(run_task_time, time_cost_str):
+        now_time = int(time.time()*1000)
+        if run_task_time == 0:
+            time_cost_str = 'start'
+        else:
+            if time_cost_str != '': 
+                time_cost_str += f'-->'
+            time_cost_str += f'{now_time - run_task_time}'
+        run_task_time = now_time
+        return run_task_time, time_cost_str
+
+
+    def enhance(
+        pil_image: Image,
+    ):
+        img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+
+        h, w = img.shape[0:2]
+        if h < 300:
+            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+        output, _ = upsampler.enhance(img, outscale=2)
+        pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+
+        return pil_output
+
+    with gr.Blocks() as demo:
+        croper = gr.State()
+        with gr.Row():
+            with gr.Column():
+                input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
+                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
+                category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
+            with gr.Column():
+                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
+                start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step")
+                with gr.Accordion("Advanced Options", open=False):
+                    guidance_scale = gr.Slider(minimum=0, maximum=20, value=1, step=0.5, label="Guidance Scale")
+                    generate_size = gr.Number(label="Generate Size", value=512)
+                    mask_expansion = gr.Number(label="Mask Expansion", value=10, visible=True)
+                    mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
+                    adapter_weights = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Adapter Weights", visible=False)
+            with gr.Column():
+                seed = gr.Number(label="Seed", value=8)
+                w1 = gr.Number(label="W1", value=2)
+                g_btn = gr.Button("Edit Image")
+                
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="pil")
+            with gr.Column():
+                restored_image = gr.Image(label="Restored Image", type="pil", interactive=False)
+                download_path = gr.File(label="Download the output image", interactive=False)
+            with gr.Column():
+                origin_area_image = gr.Image(label="Origin Area Image", type="pil", interactive=False)
+                enhanced_image = gr.Image(label="Enhanced Image", type="pil", interactive=False)
+                generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
+                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+        
+        g_btn.click(
+            fn=segment_image,
+            inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
+            outputs=[origin_area_image, croper],
+        ).success(
+            fn=image_to_image,
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, adapter_weights],
+            outputs=[enhanced_image, generated_image, generated_cost],
+        ).success(
+            fn=restore_result,
+            inputs=[croper, category, enhanced_image],
+            outputs=[restored_image, download_path],
+        )
+
+    return demo

inversion_run_base.py

@@ -12,6 +12,14 @@ from config import get_config, get_num_steps_actual
 from functools import partial
 from compel import Compel, ReturnedEmbeddingsType
 
+os.system("pip freeze")
+if not os.path.exists('GFPGANv1.4.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+os.makedirs('output', exist_ok=True)
+
 class Object(object):
     pass