Fix OOM issue.

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 import argparse
 from omegaconf import OmegaConf
-from gligen.task_grounded_generation import grounded_generation_box, load_ckpt
+from gligen.task_grounded_generation import grounded_generation_box, load_ckpt, load_common_ckpt
 
 import json
 import numpy as np
@@ -34,41 +34,46 @@ def parse_option():
 args = parse_option()
 
 
-def load_from_hf(repo_id, filename='diffusion_pytorch_model.bin'):
-    cache_file = hf_hub_download(repo_id=repo_id, filename=filename)
+def load_from_hf(repo_id, filename='diffusion_pytorch_model.bin', subfolder=None):
+    cache_file = hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder)
     return torch.load(cache_file, map_location='cpu')
 
 def load_ckpt_config_from_hf(modality):
-    ckpt = load_from_hf(f'gligen/{modality}')
-    config = load_from_hf('gligen/demo_config_legacy', filename=f'{modality}.pth')
+    ckpt = load_from_hf('gligen/demo_ckpts_legacy', filename=f'{modality}.pth', subfolder='model')
+    config = load_from_hf('gligen/demo_ckpts_legacy', filename=f'{modality}.pth', subfolder='config')
     return ckpt, config
 
 
-if args.load_text_box_generation:
-    pretrained_ckpt_gligen, config = load_ckpt_config_from_hf('gligen-generation-text-box')
+def ckpt_load_helper(modality, is_inpaint, is_style, common_instances=None):
+    pretrained_ckpt_gligen, config = load_ckpt_config_from_hf(modality)
     config = OmegaConf.create( config["_content"] ) # config used in training
     config.update( vars(args) )
-    config.model['params']['is_inpaint'] = False
-    config.model['params']['is_style'] = False
-    loaded_model_list = load_ckpt(config, pretrained_ckpt_gligen) 
+    config.model['params']['is_inpaint'] = is_inpaint
+    config.model['params']['is_style'] = is_style
 
+    if common_instances is None:
+        common_ckpt = load_from_hf('gligen/demo_ckpts_legacy', filename=f'common.pth', subfolder='model')
+        common_instances = load_common_ckpt(config, common_ckpt)
 
-if args.load_text_box_inpainting:
-    pretrained_ckpt_gligen_inpaint, config = load_ckpt_config_from_hf('gligen-inpainting-text-box')
-    config = OmegaConf.create( config["_content"] ) # config used in training
-    config.update( vars(args) )
-    config.model['params']['is_inpaint'] = True 
-    config.model['params']['is_style'] = False
-    loaded_model_list_inpaint = load_ckpt(config, pretrained_ckpt_gligen_inpaint)
+    loaded_model_list = load_ckpt(config, pretrained_ckpt_gligen, common_instances)
 
+    return loaded_model_list, common_instances
 
-if args.load_text_image_box_generation:
-    pretrained_ckpt_gligen_style, config = load_ckpt_config_from_hf('gligen-generation-text-image-box')
-    config = OmegaConf.create( config["_content"] ) # config used in training
-    config.update( vars(args) )
-    config.model['params']['is_inpaint'] = False 
-    config.model['params']['is_style'] = True
-    loaded_model_list_style = load_ckpt(config, pretrained_ckpt_gligen_style)
+
+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]
+
+loaded_model_list_style = ckpt_load_helper(
+    'gligen-generation-text-image-box',
+    is_inpaint=False, is_style=True, common_instances=common_instances
+)[0]
 
 
 def load_clip_model():

gligen/task_grounded_generation.py

@@ -65,21 +65,25 @@ def draw_box(img, locations):
         draw.rectangle([box[0]*WW, box[1]*HH, box[2]*WW, box[3]*HH], outline =colors[bid % len(colors)], width=5)
     return img 
 
-def load_ckpt(config, state_dict):
-    model = instantiate_from_config(config.model).to(device).eval()
+def load_common_ckpt(config, common_ckpt):
     autoencoder = instantiate_from_config(config.autoencoder).to(device).eval()
     text_encoder = instantiate_from_config(config.text_encoder).to(device).eval()
     diffusion = instantiate_from_config(config.diffusion).to(device)
 
-    autoencoder.load_state_dict( state_dict["autoencoder"]  )
-    text_encoder.load_state_dict( state_dict["text_encoder"]  )
-    diffusion.load_state_dict( state_dict["diffusion"]  )
+    autoencoder.load_state_dict( common_ckpt["autoencoder"]  )
+    text_encoder.load_state_dict( common_ckpt["text_encoder"]  )
+    diffusion.load_state_dict( common_ckpt["diffusion"]  )
+
+    return [autoencoder, text_encoder, diffusion]
+
+def load_ckpt(config, state_dict, common_instances):
+    model = instantiate_from_config(config.model).to(device).eval()
 
     model.load_state_dict(state_dict['model'])
     set_alpha_scale(model, config.alpha_scale)
     print("ckpt is loaded")
 
-    return model, autoencoder, text_encoder, diffusion
+    return [model] + common_instances