From 403222bfeebd751a90add67f9d877de1ddbebe57 Mon Sep 17 00:00:00 2001
From: Konie <guofengzuo@hotmail.com>
Date: Wed, 11 Oct 2023 13:08:10 +0800
Subject: [PATCH] Pre download all require model when setup

---
 fooocusapi/parameters.py | 3 +++
 fooocusapi/worker.py     | 4 ++--
 predict.py               | 8 +++++++-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/fooocusapi/parameters.py b/fooocusapi/parameters.py
index 16ad24b..a81c699 100644
--- 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
index f0fa382..510ae3f 100644
--- 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
index 1c7b1bc..ea71079 100644
--- 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(