fix: memory efficient

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from setup import setup
-import cv2
+import torch
+import gc
 from PIL import Image
 from manga_line_extraction.model import MangaLineExtractor
 from anime2sketch.model import Anime2Sketch
@@ -9,24 +10,36 @@ setup()
 
 print("Setup finished")
 
-extractor = MangaLineExtractor("./models/erika.pth", "cpu")
-to_sketch = Anime2Sketch("./models/netG.pth", "cpu")
 
-print("Model loaded")
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
 
 
 def extract(image):
-    return extractor.predict(image)
+    extractor = MangaLineExtractor("./models/erika.pth", "cpu")
+    result = extractor.predict(image)
+    del extractor
+    flush()
+    return result
 
 
 def convert_to_sketch(image):
-    return to_sketch.predict(image)
+    to_sketch = Anime2Sketch("./models/netG.pth", "cpu")
+    result = to_sketch.predict(image)
+    del to_sketch
+    flush()
+    return result
 
 
 def start(image):
     return [extract(image), convert_to_sketch(Image.fromarray(image).convert("RGB"))]
 
 
+def clear():
+    return [None, None]
+
+
 def ui():
     with gr.Blocks() as blocks:
         gr.Markdown(
@@ -43,7 +56,8 @@ def ui():
             with gr.Column():
                 input_img = gr.Image(label="Input", interactive=True)
 
-                extract_btn = gr.Button("Extract", variant="primary")
+                extract_btn = gr.Button("Start", variant="primary")
+                clear_btn = gr.Button("Clear", variant="secondary")
 
             with gr.Column():
                 # with gr.Row():
@@ -78,6 +92,12 @@ def ui():
             outputs=[extract_output_img, to_sketch_output_img],
         )
 
+        clear_btn.click(
+            fn=clear,
+            inputs=[],
+            outputs=[extract_output_img, to_sketch_output_img],
+        )
+
     return blocks