Update app.py

app.py

@@ -133,8 +133,8 @@ models_rbm.generator.eval().requires_grad_(False)
 def infer(ref_style_file, style_description, caption, use_low_vram, progress):
     global models_rbm, models_b, device
     
-    if use_low_vram:
-        models_to(models_rbm, device=device, excepts=["generator", "previewer"])
+    models_to(models_rbm, device=device)
+    
     try:
         
         caption = f"{caption} in {style_description}"
@@ -234,6 +234,8 @@ def infer(ref_style_file, style_description, caption, use_low_vram, progress):
         return sampled_image # Return the sampled_image PIL image
 
     finally:
+        if use_low_vram:
+            models_to(models_rbm, device=device)
         # Clear CUDA cache
         torch.cuda.empty_cache()
         gc.collect()
@@ -241,10 +243,9 @@ def infer(ref_style_file, style_description, caption, use_low_vram, progress):
 def infer_compo(style_description, ref_style_file, caption, ref_sub_file, use_low_vram, progress):
     global models_rbm, models_b, device
     sam_model = LangSAM()
-    if use_low_vram:
-        models_to(models_rbm, device=device, excepts=["generator", "previewer"])
-        models_to(sam_model, device=device)
-        models_to(sam_model.sam, device=device)
+    models_to(models_rbm, device=device)
+    models_to(sam_model, device=device)
+    models_to(sam_model.sam, device=device)
     try:
         caption = f"{caption} in {style_description}"
         sam_prompt = f"{caption}"
@@ -361,6 +362,10 @@ def infer_compo(style_description, ref_style_file, caption, ref_sub_file, use_lo
         return sampled_image  # Return the sampled_image PIL image
 
     finally:
+        if use_low_vram:
+            models_to(models_rbm, device=device, excepts=["generator", "previewer"])
+            models_to(sam_model, device=device)
+            models_to(sam_model.sam, device=device)
         # Clear CUDA cache
         torch.cuda.empty_cache()
         gc.collect()