Update app.py

app.py

@@ -107,28 +107,32 @@ models_b = WurstCoreB.Models(
 )
 models_b.generator.bfloat16().eval().requires_grad_(False)
 
-if low_vram:
-    # Off-load old generator (which is not used in models_rbm)
-    models.generator.to("cpu")
-    torch.cuda.empty_cache()
 
-generator_rbm = StageCRBM()
-for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items():
-    set_module_tensor_to_device(generator_rbm, param_name, "cpu", value=param)
-generator_rbm = generator_rbm.to(getattr(torch, core.config.dtype)).to(device)
-generator_rbm = core.load_model(generator_rbm, 'generator')
-
-models_rbm = core.Models(
-        effnet=models.effnet, previewer=models.previewer,
-        generator=generator_rbm, generator_ema=models.generator_ema,
-        tokenizer=models.tokenizer, text_model=models.text_model, image_model=models.image_model
-    )
-models_rbm.generator.eval().requires_grad_(False)
 
 sam_model = LangSAM()
 
 def infer(ref_style_file, style_description, caption, progress):
     global models_rbm, models_b, device
+
+    if low_vram:
+        # Off-load old generator (which is not used in models_rbm)
+        models.generator.to("cpu")
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    generator_rbm = StageCRBM()
+    for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items():
+        set_module_tensor_to_device(generator_rbm, param_name, "cpu", value=param)
+    generator_rbm = generator_rbm.to(getattr(torch, core.config.dtype)).to(device)
+    generator_rbm = core.load_model(generator_rbm, 'generator')
+
+    models_rbm = core.Models(
+        effnet=models.effnet, previewer=models.previewer,
+        generator=generator_rbm, generator_ema=models.generator_ema,
+        tokenizer=models.tokenizer, text_model=models.text_model, image_model=models.image_model
+    )
+    models_rbm.generator.eval().requires_grad_(False)
+    
     if low_vram:
         models_to(models_rbm, device=device, excepts=["generator", "previewer"])
     try:
@@ -236,6 +240,26 @@ def infer(ref_style_file, style_description, caption, progress):
 
 def infer_compo(style_description, ref_style_file, caption, ref_sub_file, progress):
     global models_rbm, models_b, device, sam_model
+
+    if low_vram:
+        # Off-load old generator (which is not used in models_rbm)
+        models.generator.to("cpu")
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    generator_rbm = StageCRBM()
+    for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items():
+        set_module_tensor_to_device(generator_rbm, param_name, "cpu", value=param)
+    generator_rbm = generator_rbm.to(getattr(torch, core.config.dtype)).to(device)
+    generator_rbm = core.load_model(generator_rbm, 'generator')
+
+    models_rbm = core.Models(
+        effnet=models.effnet, previewer=models.previewer,
+        generator=generator_rbm, generator_ema=models.generator_ema,
+        tokenizer=models.tokenizer, text_model=models.text_model, image_model=models.image_model
+    )
+    models_rbm.generator.eval().requires_grad_(False)
+    
     if low_vram:
         models_to(models_rbm, device=device, excepts=["generator", "previewer"])
         models_to(sam_model, device=device)