Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import contextlib
 import gc
 import json
@@ -179,6 +180,7 @@ image_encoder = CLIPVisionModelWithProjection.from_pretrained(os.path.join(model
 global pipeline
 global MultiResNetModel
 
+@spaces.GPU
 def load_ckpt():
     global pipeline
     global MultiResNetModel
@@ -293,6 +295,7 @@ def load_ckpt():
     
 global cur_style
 cur_style = 'line + shadow'
+@spaces.GPU
 def change_ckpt(style):
     global pipeline
     global MultiResNetModel
@@ -334,6 +337,7 @@ def change_ckpt(style):
   
 load_ckpt()
 
+@spaces.GPU
 def fix_random_seeds(seed):
     random.seed(seed)
     np.random.seed(seed)
@@ -349,6 +353,7 @@ def process_multi_images(files):
         imgs.append(img)
     return imgs 
 
+@spaces.GPU
 def extract_lines(image):
     src = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
 
@@ -373,16 +378,17 @@ def extract_lines(image):
     torch.cuda.empty_cache()
     return outimg
 
+@spaces.GPU
 def extract_line_image(query_image_, resolution):
     tar_width, tar_height = resolution
     query_image = query_image_.resize((tar_width, tar_height))
-    # query_image.save('/mnt/workspace/zhuangjunhao/cobra_code/ColorFlow/examples/line/example3/input.png')
     query_image = query_image.convert('L').convert('RGB')
     extracted_line = extract_lines(query_image)
     extracted_line = extracted_line.convert('L').convert('RGB')
     torch.cuda.empty_cache()
     return extracted_line, Image.new('RGB', (tar_width, tar_height), 'black')
 
+@spaces.GPU
 def extract_sketch_line_image(query_image_, input_style):
     global cur_style
     if input_style != cur_style:
@@ -418,6 +424,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):
     if extracted_line is None:
         gr.Info("Please preprocess the image first")
@@ -440,11 +447,6 @@ def colorize_image(extracted_line, reference_images, resolution, seed, num_infer
     reference_images = [process_image(ref_image, tar_width, tar_height) for ref_image in reference_images]
     query_patches_pil = process_image_Q_varres(query_image_origin, tar_width, tar_height)
     reference_patches_pil = []
-    # Save reference_images
-    # save_path = '/mnt/workspace/zhuangjunhao/cobra_code/ColorFlow/examples/line/example3'
-    # os.makedirs(save_path, exist_ok=True)
-    # for idx, ref_image in enumerate(reference_images):
-    #     ref_image.save(os.path.join(save_path, f'reference_image_{idx}.png'))
 
     for reference_image in reference_images:
         reference_patches_pil += process_image_ref_varres(reference_image, tar_width, tar_height)
@@ -695,4 +697,4 @@ with gr.Blocks() as demo:
         )
 
 
-demo.launch(server_name="0.0.0.0", server_port=52218)
+demo.launch()