Pre download all require model when setup

mrhan1993 · Oct 11, 2023 · 403222b · 403222b
1 parent b7f496b
commit 403222b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/fooocusapi/parameters.py b/fooocusapi/parameters.py
@@ -3,6 +3,9 @@
 import numpy as np
 
 
+inpaint_model_version = 'v1'
+
+
 fooocus_styles = [
     'Fooocus V2',
     'Default (Slightly Cinematic)',

diff --git a/fooocusapi/worker.py b/fooocusapi/worker.py
@@ -4,7 +4,7 @@
 import numpy as np
 import torch
 from typing import List
-from fooocusapi.parameters import GenerationFinishReason, ImageGenerationParams, ImageGenerationResult
+from fooocusapi.parameters import inpaint_model_version, GenerationFinishReason, ImageGenerationParams, ImageGenerationResult
 from fooocusapi.task_queue import TaskQueue, TaskType
 
 save_log = True
@@ -118,7 +118,7 @@ def build_advanced_parameters():
             controlnet_softness = 0.25
             canny_low_threshold = 64
             canny_high_threshold = 128
-            inpaint_engine = 'v1'
+            inpaint_engine = inpaint_model_version
             freeu_enabled = False
             freeu_b1, freeu_b2, freeu_s1, freeu_s2 = [None] * 4
             return [adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adaptive_cfg, sampler_name,

diff --git a/predict.py b/predict.py
@@ -6,7 +6,7 @@
 from typing import List
 from cog import BasePredictor, Input, Path
 
-from fooocusapi.parameters import GenerationFinishReason, ImageGenerationParams, fooocus_styles, aspect_ratios, uov_methods, outpaint_expansions
+from fooocusapi.parameters import inpaint_model_version, GenerationFinishReason, ImageGenerationParams, fooocus_styles, aspect_ratios, uov_methods, outpaint_expansions
 from fooocusapi.worker import process_generate
 import numpy as np
 from PIL import Image
@@ -32,6 +32,12 @@ def setup(self) -> None:
 
         print("[Predictor Setup] Preload pipeline")
         import modules.default_pipeline as _
+        import modules.path as path
+        path.downloading_upscale_model()
+        path.downloading_inpaint_models(inpaint_model_version)
+        path.downloading_controlnet_canny()
+        path.downloading_controlnet_cpds()
+        path.downloading_ip_adapters()
         print("[Predictor Setup] Finished")
 
     def predict(