model2 commited on
Commit
3d636e1
·
1 Parent(s): 04dde6a

fix preloading

Browse files
Files changed (1) hide show
  1. app.py +10 -13
app.py CHANGED
@@ -79,29 +79,27 @@ def import_custom_nodes() -> None:
79
 
80
 
81
  # Preload nodes, models.
 
82
  loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
83
  upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
84
  reactorloadfacemodel = NODE_CLASS_MAPPINGS["ReActorLoadFaceModel"]()
85
- face_model = reactorloadfacemodel.load_model(
86
  face_model="advance_face_model.safetensors"
87
  )
88
  imageresize = NODE_CLASS_MAPPINGS["ImageResize+"]()
89
  reactorfaceswap = NODE_CLASS_MAPPINGS["ReActorFaceSwap"]()
90
  imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
91
  saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
92
- upscale_model = upscalemodelloader.load_model(model_name="ESRGAN/4x_NMKD-Siax_200k.pth")
93
  model_loaders = [
94
- upscalemodelloader,
95
- reactorloadfacemodel,
96
  reactorfaceswap,
97
  imageupscalewithmodel,
98
  ]
99
- valid_models = [
100
- getattr(loader[0], "patcher", loader[0])
101
- for loader in model_loaders
102
- if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
103
- ]
104
- model_management.load_models_gpu(valid_models)
105
 
106
 
107
  def load_extra_path_config(yaml_path):
@@ -214,7 +212,6 @@ add_extra_model_paths()
214
 
215
  @spaces.GPU(duration=60)
216
  def advance_blur(input_image):
217
- import_custom_nodes()
218
  with torch.inference_mode():
219
  loaded_input_image = loadimage.load_image(
220
  image=input_image,
@@ -243,11 +240,11 @@ def advance_blur(input_image):
243
  source_faces_index="0",
244
  console_log_level=2,
245
  input_image=get_value_at_index(resized_input_image, 0),
246
- face_model=get_value_at_index(face_model, 0),
247
  )
248
 
249
  upscaled_image = imageupscalewithmodel.upscale(
250
- upscale_model=get_value_at_index(upscale_model, 0),
251
  image=get_value_at_index(swapped_image, 0),
252
  )
253
 
 
79
 
80
 
81
  # Preload nodes, models.
82
+ import_custom_nodes()
83
  loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
84
  upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
85
  reactorloadfacemodel = NODE_CLASS_MAPPINGS["ReActorLoadFaceModel"]()
86
+ FACE_MODEL = reactorloadfacemodel.load_model(
87
  face_model="advance_face_model.safetensors"
88
  )
89
  imageresize = NODE_CLASS_MAPPINGS["ImageResize+"]()
90
  reactorfaceswap = NODE_CLASS_MAPPINGS["ReActorFaceSwap"]()
91
  imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
92
  saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
93
+ UPSCALE_MODEL = upscalemodelloader.load_model(model_name="ESRGAN/4x_NMKD-Siax_200k.pth")
94
  model_loaders = [
95
+ UPSCALE_MODEL,
96
+ FACE_MODEL,
97
  reactorfaceswap,
98
  imageupscalewithmodel,
99
  ]
100
+ model_management.load_models_gpu([
101
+ loader[0].patcher if hasattr(loader[0], 'patcher') else loader[0] for loader in model_loaders
102
+ ])
 
 
 
103
 
104
 
105
  def load_extra_path_config(yaml_path):
 
212
 
213
  @spaces.GPU(duration=60)
214
  def advance_blur(input_image):
 
215
  with torch.inference_mode():
216
  loaded_input_image = loadimage.load_image(
217
  image=input_image,
 
240
  source_faces_index="0",
241
  console_log_level=2,
242
  input_image=get_value_at_index(resized_input_image, 0),
243
+ face_model=get_value_at_index(FACE_MODEL, 0),
244
  )
245
 
246
  upscaled_image = imageupscalewithmodel.upscale(
247
+ upscale_model=get_value_at_index(UPSCALE_MODEL, 0),
248
  image=get_value_at_index(swapped_image, 0),
249
  )
250