Update app.py

app.py

@@ -382,9 +382,6 @@ def extract_line_image(query_image_, resolution):
 
 @spaces.GPU
 def extract_sketch_line_image(query_image_, input_style):
-    global cur_style
-    if input_style != cur_style:
-        change_ckpt(input_style)
 
     resolution = get_rate(query_image_)
     extracted_line, hint_mask = extract_line_image(query_image_, resolution)
@@ -417,7 +414,7 @@ def extract_sketch_line_image(query_image_, input_style):
     return extracted_sketch_line.convert('RGB'), extracted_sketch_line.convert('RGB'), hint_mask, query_image_, extracted_sketch_line_ori.convert('RGB'), resolution
 
 @spaces.GPU(duration=120)
-def colorize_image(extracted_line, reference_images, resolution, seed, num_inference_steps, top_k, hint_mask=None, hint_color=None, query_image_origin=None, extracted_image_ori=None):
+def colorize_image(input_style, extracted_line, reference_images, resolution, seed, num_inference_steps, top_k, hint_mask=None, hint_color=None, query_image_origin=None, extracted_image_ori=None):
     if extracted_line is None:
         gr.Info("Please preprocess the image first")
         raise ValueError("Please preprocess the image first")
@@ -427,6 +424,8 @@ def colorize_image(extracted_line, reference_images, resolution, seed, num_infer
     global pipeline
     global MultiResNetModel
     global cur_style
+    if input_style != cur_style:
+        change_ckpt(input_style)
 
     tar_width, tar_height = resolution
 
@@ -678,7 +677,7 @@ with gr.Blocks() as demo:
     )
     colorize_button.click(
         colorize_image, 
-        inputs=[extracted_image, reference_images, resolution, seed, num_inference_steps, top_k, hint_mask, hint_color, query_image_origin, extracted_image_ori], 
+        inputs=[model_name, extracted_image, reference_images, resolution, seed, num_inference_steps, top_k, hint_mask, hint_color, query_image_origin, extracted_image_ori], 
         outputs=output_gallery
     )
     with gr.Column():