Update app.py

app.py

@@ -152,6 +152,22 @@ examples = [
         0,  
         10  
     ],
+    [
+        "./assets/example_6/input.png", 
+        ["./assets/example_6/ref1.png", "./assets/example_6/ref2.png", "./assets/example_6/ref3.png"], 
+        "Sketch_Shading", 
+        "512x800",  
+        0, 
+        10 
+    ],
+    [
+        "./assets/example_7/input.png", 
+        ["./assets/example_7/ref1.png", "./assets/example_7/ref2.png", "./assets/example_7/ref3.png"， "./assets/example_7/ref4.png"], 
+        "Sketch_Shading", 
+        "640x640",  
+        2, 
+        10 
+    ],
     [
         "./assets/example_1/input.jpg", 
         ["./assets/example_1/ref1.jpg", "./assets/example_1/ref2.jpg", "./assets/example_1/ref3.jpg"], 
@@ -177,9 +193,13 @@ global MultiResNetModel
 def load_ckpt(input_style):
     global pipeline
     global MultiResNetModel
-    if input_style == "Sketch":
-        ckpt_path = model_global_path + '/sketch/'
-        rank = 128
+    if input_style == "Sketch" or input_style == "Sketch_Shading":
+        if input_style == "Sketch":
+            ckpt_path = './ckpt/sketch/'
+            rank = 128
+        else:
+            ckpt_path = './ckpt/shading/'
+            rank = 128
         pretrained_model_name_or_path = 'PixArt-alpha/PixArt-XL-2-1024-MS'
         transformer = PixArtTransformer2DModel.from_pretrained(
             pretrained_model_name_or_path, subfolder="transformer", revision=None, variant=None
@@ -345,6 +365,17 @@ def extract_line_image(query_image_, input_style, resolution):
         extracted_line = extract_lines(query_image)
         extracted_line = extracted_line.convert('L').convert('RGB')
         input_context = extracted_line
+    elif input_style == "Sketch_Shading":
+        query_image = query_image.convert('L').convert('RGB')
+        extracted_line = extract_lines(query_image)
+        extracted_line = extracted_line.convert('L').convert('RGB')
+        array1 = np.array(query_image)
+        array2 = np.array(extracted_line)
+        array2[array1 < 0.3 * 255.0] = 0
+        gray_rate = 125
+        up_bound = 145
+        array2[(array2 > gray_rate) & (array1 < up_bound) & (array1 > 0.3 * 255.0)] = gray_rate
+        input_context = Image.fromarray(np.uint8(array2))
     torch.cuda.empty_cache()
     return input_context, extracted_line, input_context  
 
@@ -461,7 +492,7 @@ with gr.Blocks() as demo:
 </div>
 <div style="text-align: left; margin: 0 auto;">
     <ol style="font-size: 1.1em;">
-        <li>Choose input style: GrayImage(ScreenStyle) or Sketch.</li>
+        <li>Choose input style: GrayImage(ScreenStyle)、Sketch with Shading or Sketch.</li>
         <li>Upload your image: Use the 'Upload' button to select the image you want to colorize.</li>
         <li>Preprocess the image: Click the 'Preprocess' button to decolorize the image.</li>
         <li>Upload reference images: Upload multiple reference images to guide the colorization.</li>
@@ -481,7 +512,7 @@ with gr.Blocks() as demo:
 </div>
 <div style="text-align: left; margin: 0 auto;">
     <ol style="font-size: 1.1em;">
-        <li>选择输入样式：灰度图(ScreenStyle)、线稿。</li>
+        <li>选择输入样式：灰度图(ScreenStyle)、线稿+阴影、线稿。</li>
         <li>上传您的图像：使用“上传”按钮选择要上色的图像。</li>
         <li>预处理图像：点击“预处理”按钮以去色图像。</li>
         <li>上传参考图像：上传多张参考图像以指导上色。</li>
@@ -499,7 +530,7 @@ with gr.Blocks() as demo:
     
     with gr.Column():
         with gr.Row():
-            input_style = gr.Radio(["GrayImage(ScreenStyle)", "Sketch"], label="Input Style", value="GrayImage(ScreenStyle)")
+            input_style = gr.Radio(["GrayImage(ScreenStyle)", "Sketch_Shading", "Sketch"], label="Input Style", value="GrayImage(ScreenStyle)")
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(type="pil", label="Image to Colorize")