Update

app.py

@@ -7,7 +7,9 @@ import json
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 from functools import partial
+from collections import Counter
 import math
+import gc
 
 from gradio import processing_utils
 from typing import Optional
@@ -42,20 +44,56 @@ def ckpt_load_helper(modality, is_inpaint, is_style, common_instances=None):
     return loaded_model_list, common_instances
 
 
-loaded_model_list, common_instances = ckpt_load_helper(
-    'gligen-generation-text-box',
-    is_inpaint=False, is_style=False, common_instances=None
-)
-
-loaded_model_list_inpaint = ckpt_load_helper(
-    'gligen-inpainting-text-box',
-    is_inpaint=True, is_style=False, common_instances=common_instances
-)[0]
+class Instance:
+    def __init__(self, capacity = 2):
+        self.model_type = 'base'
+        self.loaded_model_list = {}
+        self.counter = Counter()
+        self.counter['base'] = 0
+        self.loaded_model_list['base'], self.common_instances = ckpt_load_helper(
+            'gligen-generation-text-box',
+            is_inpaint=False, is_style=False, common_instances=None
+        )
+        self.capacity = capacity
+
+    def get_model(self, model_type):
+        if model_type in self.loaded_model_list:
+            self.counter[model_type] += 1
+            print(self.counter)
+            return self.loaded_model_list[model_type]
+
+        if self.capacity == len(self.loaded_model_list):
+            least_used_type = self.counter.most_common()[-1][0]
+            del self.loaded_model_list[least_used_type]
+            del self.counter[least_used_type]
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        self.counter[model_type] = 1
+        self.loaded_model_list[model_type] = self._get_model(model_type)
+        print(self.counter)
+        return self.loaded_model_list[model_type]
+
+    def _get_model(self, model_type):
+        if model_type == 'base':
+            return ckpt_load_helper(
+                'gligen-generation-text-box',
+                is_inpaint=False, is_style=False, common_instances=self.common_instances
+            )[0]
+        elif model_type == 'inpaint':
+            return ckpt_load_helper(
+                'gligen-inpainting-text-box',
+                is_inpaint=True, is_style=False, common_instances=self.common_instances
+            )[0]
+        elif model_type == 'style':
+            return ckpt_load_helper(
+                'gligen-generation-text-image-box',
+                is_inpaint=False, is_style=True, common_instances=self.common_instances
+            )[0]
+        
+        assert False
 
-loaded_model_list_style = ckpt_load_helper(
-    'gligen-generation-text-image-box',
-    is_inpaint=False, is_style=True, common_instances=common_instances
-)[0]
+instance = Instance()
 
 
 def load_clip_model():
@@ -143,7 +181,7 @@ def inference(task, language_instruction, grounding_instruction, inpainting_boxe
     image_list = [placeholder_image] * len(phrase_list) # placeholder input for visual prompt, which is disabled
 
     batch_size = int(batch_size)
-    if not 1 <= batch_size <= 2:
+    if not 1 <= batch_size <= 4:
         batch_size = 2
 
     if style_image == None:
@@ -183,13 +221,13 @@ def inference(task, language_instruction, grounding_instruction, inpainting_boxe
     with torch.autocast(device_type='cuda', dtype=torch.float16):
         if task == 'Grounded Generation':
             if style_image == None:
-                return grounded_generation_box(loaded_model_list, instruction, *args, **kwargs)
+                return grounded_generation_box(instance.get_model('base'), instruction, *args, **kwargs)
             else:
-                return grounded_generation_box(loaded_model_list_style, instruction, *args, **kwargs)
+                return grounded_generation_box(instance.get_model('style'), instruction, *args, **kwargs)
         elif task == 'Grounded Inpainting':
             assert image is not None
             instruction['input_image'] = image.convert("RGB")
-            return grounded_generation_box(loaded_model_list_inpaint, instruction, *args, **kwargs)
+            return grounded_generation_box(instance.get_model('inpaint'), instruction, *args, **kwargs)
 
 
 def draw_box(boxes=[], texts=[], img=None):
@@ -498,7 +536,7 @@ with Blocks(
                 with gr.Column():
                     alpha_sample = gr.Slider(minimum=0, maximum=1.0, step=0.1, value=0.3, label="Scheduled Sampling (τ)")
                     guidance_scale = gr.Slider(minimum=0, maximum=50, step=0.5, value=7.5, label="Guidance Scale")
-                    batch_size = gr.Slider(minimum=1, maximum=2, step=1, value=2, label="Number of Samples")
+                    batch_size = gr.Slider(minimum=1, maximum=4, step=1, value=2, label="Number of Samples")
                     append_grounding = gr.Checkbox(value=True, label="Append grounding instructions to the caption")
                     use_actual_mask = gr.Checkbox(value=False, label="Use actual mask for inpainting", visible=False)
                     with gr.Row():