add use low vram option

app.py

@@ -34,7 +34,7 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 print(device)
 
 # Flag for low VRAM usage
-low_vram = False
+# low_vram = False
 
 # Function definition for low VRAM usage
 def models_to(model, device="cpu", excepts=None):
@@ -107,11 +107,13 @@ 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()
     gc.collect()
+"""
 
 generator_rbm = StageCRBM()
 for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items():
@@ -128,10 +130,10 @@ models_rbm.generator.eval().requires_grad_(False)
 
 
 
-def infer(ref_style_file, style_description, caption, progress):
+def infer(ref_style_file, style_description, caption, use_low_vram, progress):
     global models_rbm, models_b, device
     
-    if low_vram:
+    if use_low_vram:
         models_to(models_rbm, device=device, excepts=["generator", "previewer"])
     try:
         
@@ -167,7 +169,7 @@ def infer(ref_style_file, style_description, caption, progress):
         conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
         unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
 
-        if low_vram:
+        if use_low_vram:
             # The sampling process uses more vram, so we offload everything except two modules to the cpu.
             models_to(models_rbm, device="cpu", excepts=["generator", "previewer"])
 
@@ -236,10 +238,10 @@ def infer(ref_style_file, style_description, caption, progress):
         torch.cuda.empty_cache()
         gc.collect()
 
-def infer_compo(style_description, ref_style_file, caption, ref_sub_file, 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 low_vram:
+    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)
@@ -288,7 +290,7 @@ def infer_compo(style_description, ref_style_file, caption, ref_sub_file, progre
         conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
         unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
 
-        if low_vram:
+        if use_low_vram:
             models_to(models_rbm, device="cpu", excepts=["generator", "previewer"])
             models_to(sam_model, device="cpu")
             models_to(sam_model.sam, device="cpu")
@@ -363,13 +365,13 @@ def infer_compo(style_description, ref_style_file, caption, ref_sub_file, progre
         torch.cuda.empty_cache()
         gc.collect()
 
-def run(style_reference_image, style_description, subject_prompt, subject_reference, use_subject_ref):
+def run(style_reference_image, style_description, subject_prompt, subject_reference, use_subject_ref, use_low_vram):
     result = None
     progress = gr.Progress(track_tqdm=True)
     if use_subject_ref is True:
-        result = infer_compo(style_description, style_reference_image, subject_prompt, subject_reference, progress)
+        result = infer_compo(style_description, style_reference_image, subject_prompt, subject_reference, use_low_vram, progress)
     else:
-        result = infer(style_reference_image, style_description, subject_prompt, progress)
+        result = infer(style_reference_image, style_description, subject_prompt, use_low_vram, progress)
     return result
 
 def show_hide_subject_image_component(use_subject_ref):
@@ -406,7 +408,9 @@ with gr.Blocks(analytics_enabled=False) as demo:
                 subject_prompt = gr.Textbox(
                     label = "Subject Prompt"
                 )
-                use_subject_ref = gr.Checkbox(label="Use Subject Image as Reference", value=False)
+                with gr.Row():
+                    use_subject_ref = gr.Checkbox(label="Use Subject Image as Reference", value=False)
+                    use_low_vram = gr.Checkbox(label="Use Low-VRAM", value=False)
                 
                 with gr.Accordion("Advanced Settings", open=False) as sub_img_panel:
                     subject_reference = gr.Image(label="Subject Reference", type="filepath")
@@ -418,13 +422,13 @@ with gr.Blocks(analytics_enabled=False) as demo:
                 output_image = gr.Image(label="Output Image")
                 gr.Examples(
                     examples = [
-                        ["./data/cyberpunk.png", "cyberpunk art style", "a car", None, False],
-                        ["./data/mosaic.png", "mosaic art style", "a lighthouse", None, False],
-                        ["./data/glowing.png", "glowing style", "a dwarf", None, False],
-                        ["./data/melting_gold.png", "melting golden 3D rendering style", "a dog", "./data/dog.jpg", True]
+                        ["./data/cyberpunk.png", "cyberpunk art style", "a car", None, False, False],
+                        ["./data/mosaic.png", "mosaic art style", "a lighthouse", None, False, False],
+                        ["./data/glowing.png", "glowing style", "a dwarf", None, False, False],
+                        ["./data/melting_gold.png", "melting golden 3D rendering style", "a dog", "./data/dog.jpg", True, False]
                     ],
                     fn=run,
-                    inputs=[style_reference_image, style_description, subject_prompt, subject_reference, use_subject_ref],
+                    inputs=[style_reference_image, style_description, subject_prompt, subject_reference, use_subject_ref, use_low_vram],
                     outputs=[output_image],
                     cache_examples=False
                 
@@ -439,7 +443,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     
     submit_btn.click(
         fn = run,
-        inputs = [style_reference_image, style_description, subject_prompt, subject_reference, use_subject_ref],
+        inputs = [style_reference_image, style_description, subject_prompt, subject_reference, use_subject_ref, use_low_vram],
         outputs = [output_image],
         show_api = False
     )