optimize setup

gradio_app.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-
 import gradio as gr
+import nltk
 import numpy as np
 from PIL import Image
 
-import nltk
 nltk.download('punkt')
 nltk.download('averaged_perceptron_tagger')
 
-from main import LPMConfig, main
+from main import LPMConfig, main, setup
 
 DESCRIPTION = '''# Localizing Object-level Shape Variations with Text-to-Image Diffusion Models
 This is a demo for our ''Localizing Object-level Shape Variations with Text-to-Image Diffusion Models'' [paper](https://arxiv.org/abs/2303.11306).
@@ -17,6 +16,8 @@ We introduce a method that generates object-level shape variation for a given im
 This demo allows using a real image as well as a generated image. For a real image, a matching prompt is required.
 '''
 
+stable, stable_config = setup(LPMConfig())
+
 def main_pipeline(
         prompt: str,
         object_of_interest: str,
@@ -47,7 +48,7 @@ def main_pipeline(
             real_image_path="" if input_image is None else input_image
         )
 
-        result_images, result_proxy_words = main(args)
+        result_images, result_proxy_words = main(stable, stable_config, args)
         result_images = [im.permute(1, 2, 0).cpu().numpy() for im in result_images]
         result_images = [(im * 255).astype(np.uint8) for im in result_images]
         result_images = [Image.fromarray(im) for im in result_images]

main.py

@@ -1,21 +1,20 @@
 import json
 import os
-from dataclasses import dataclass, field
-from typing import List
-
 import pyrallis
 import torch
+from dataclasses import dataclass, field
 from torch.utils.data import DataLoader
-from torchvision.utils import save_image
 from torchvision.transforms import ToTensor
+from torchvision.utils import save_image
 from tqdm import tqdm
+from typing import List
 
-from src.prompt_to_prompt_controllers import AttentionStore, AttentionReplace
-from src.null_text_inversion import invert_image
-from src.prompt_utils import get_proxy_prompts
-from src.prompt_mixing import PromptMixing
 from src.diffusion_model_wrapper import DiffusionModelWrapper, get_stable_diffusion_model, get_stable_diffusion_config, \
     generate_original_image
+from src.null_text_inversion import invert_image
+from src.prompt_mixing import PromptMixing
+from src.prompt_to_prompt_controllers import AttentionStore, AttentionReplace
+from src.prompt_utils import get_proxy_prompts
 
 
 def save_args_dict(args, similar_words):
@@ -29,10 +28,13 @@ def save_args_dict(args, similar_words):
 
     return exp_path
 
-
-def main(args):
+def setup(args):
     ldm_stable = get_stable_diffusion_model(args)
     ldm_stable_config = get_stable_diffusion_config(args)
+    return ldm_stable, ldm_stable_config
+
+
+def main(ldm_stable, ldm_stable_config, args):
 
     similar_words, prompts, another_prompts = get_proxy_prompts(args, ldm_stable)
     exp_path = save_args_dict(args, similar_words)
@@ -147,4 +149,5 @@ if __name__ == '__main__':
     args = pyrallis.parse(config_class=LPMConfig)
 
     print(args)
-    main(args)
+    stable, stable_config = setup(args)
+    main(stable, stable_config, args)