add cache_ckpt option

.gitignore

@@ -1,2 +1,3 @@
 trash/*
-tmp
+tmp
+gradio_cached_examples

app_fatezero.py

@@ -33,6 +33,7 @@ pipe = merge_config_then_run()
 
 with gr.Blocks(css='style.css') as demo:
     # gr.Markdown(TITLE)
+    
     gr.HTML(
     """
     <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
@@ -291,7 +292,9 @@ with gr.Blocks(css='style.css') as demo:
                     ],
                     outputs=result,
                     fn=pipe.run,
-                    cache_examples=os.getenv('SYSTEM') == 'spaces')
+                    cache_examples=True,
+                    # cache_examples=os.getenv('SYSTEM') == 'spaces'
+                    )
 
     # model_id.change(fn=app.load_model_info,
     #                 inputs=model_id,

inference_fatezero.py

@@ -8,32 +8,39 @@ class merge_config_then_run():
     def __init__(self) -> None:
             # Load the tokenizer
         pretrained_model_path = 'FateZero/ckpt/stable-diffusion-v1-4'
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            pretrained_model_path,
-            # 'FateZero/ckpt/stable-diffusion-v1-4',
-            subfolder="tokenizer",
-            use_fast=False,
-        )
+        self.tokenizer = None
+        self.text_encoder = None
+        self.vae = None
+        self.unet = None
+        
+        cache_ckpt = True
+        if cache_ckpt:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                pretrained_model_path,
+                # 'FateZero/ckpt/stable-diffusion-v1-4',
+                subfolder="tokenizer",
+                use_fast=False,
+            )
 
-        # Load models and create wrapper for stable diffusion
-        self.text_encoder = CLIPTextModel.from_pretrained(
-            pretrained_model_path,
-            subfolder="text_encoder",
-        )
+            # Load models and create wrapper for stable diffusion
+            self.text_encoder = CLIPTextModel.from_pretrained(
+                pretrained_model_path,
+                subfolder="text_encoder",
+            )
 
-        self.vae = AutoencoderKL.from_pretrained(
-            pretrained_model_path,
-            subfolder="vae",
-        )
-        model_config = {
-            "lora": 160,
-            # temporal_downsample_time: 4
-            "SparseCausalAttention_index": ['mid'],
-            "least_sc_channel": 640
-        }
-        self.unet = UNetPseudo3DConditionModel.from_2d_model(
-            os.path.join(pretrained_model_path, "unet"), model_config=model_config
-        )
+            self.vae = AutoencoderKL.from_pretrained(
+                pretrained_model_path,
+                subfolder="vae",
+            )
+            model_config = {
+                "lora": 160,
+                # temporal_downsample_time: 4
+                "SparseCausalAttention_index": ['mid'],
+                "least_sc_channel": 640
+            }
+            self.unet = UNetPseudo3DConditionModel.from_2d_model(
+                os.path.join(pretrained_model_path, "unet"), model_config=model_config
+            )
 
     def run(
         self,