add app_haircolor_pix2pix

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 
 # from app_base import create_demo as create_demo_face
-from app_haircolor import create_demo as create_demo_haircolor
+from app_haircolor_pix2pix import create_demo as create_demo_haircolor
 
 with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():

app_haircolor_pix2pix.py

@@ -0,0 +1,117 @@
+import spaces
+import gradio as gr
+import time
+import torch
+
+from PIL import Image
+from segment_utils import(
+    segment_image,
+    restore_result,
+)
+from diffusers import (
+    StableDiffusionInstructPix2PixPipeline,
+    EulerAncestralDiscreteScheduler,
+)
+
+BASE_MODEL = "timbrooks/instruct-pix2pix"
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+DEFAULT_EDIT_PROMPT = "change hair to blue"
+
+DEFAULT_CATEGORY = "hair"
+
+basepipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+)
+
+basepipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(basepipeline.scheduler.config)
+
+basepipeline = basepipeline.to(DEVICE)
+
+basepipeline.enable_model_cpu_offload()
+
+@spaces.GPU(duration=15)
+def image_to_image(
+    input_image: Image,
+    edit_prompt: str,
+    seed: int,
+    num_steps: int,
+    guidance_scale: float,
+    image_guidance_scale: float,
+):
+    run_task_time = 0
+    time_cost_str = ''
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
+    generated_image = basepipeline(
+        generator=generator,
+        prompt=edit_prompt,
+        image=input_image,
+        guidance_scale=guidance_scale,
+        image_guidance_scale=image_guidance_scale,
+        num_inference_steps=num_steps,
+    ).images[0]
+    
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+
+    return generated_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 create_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        croper = gr.State()
+        with gr.Row():
+            with gr.Column():
+                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
+                generate_size = gr.Number(label="Generate Size", value=512)
+            with gr.Column():
+                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
+                guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale")
+            with gr.Column():
+                image_guidance_scale = gr.Slider(minimum=0, maximum=30, value=1.5, step=0.1, label="Image Guidance Scale")
+                with gr.Accordion("Advanced Options", open=False):
+                    mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
+                    mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
+                    seed = gr.Number(label="Seed", value=8)
+                    category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
+                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)
+            with gr.Column():
+                origin_area_image = gr.Image(label="Origin Area Image", type="pil", interactive=False)
+                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+                generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, 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, edit_prompt,seed, num_steps, guidance_scale, image_guidance_scale],
+            outputs=[generated_image, generated_cost],
+        ).success(
+            fn=restore_result,
+            inputs=[croper, category, generated_image],
+            outputs=[restored_image],
+        )
+
+    return demo