Update app.py

app.py

@@ -168,7 +168,10 @@ def load_and_prepare_model():
     
     #pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1)
     #pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained('SG161222/RealVisXL_V5.0', subfolder='scheduler', algorithm_type='sde-dpmsolver++')
+    #****  NEVER TO BFLOAT BEFORE CUDA****#
     pipe.vae = vaeX #.to(torch.bfloat16)
+    #****  NEVER TO BFLOAT BEFORE CUDA****#
+
     #pipe.unet = unetX
     
     #pipe.vae.do_resize=False
@@ -201,9 +204,12 @@ def load_and_prepare_model():
     #pipe.set_adapters(["skin", "photo", "fantasy"], adapter_weights=[0.75, 0.25, 0.5])
     #pipe.set_adapters(["skin"], adapter_weights=[0.5])    
     #pipe.unet.set_default_attn_processor()
-    #pipe.to(device, torch.bfloat16)
-    pipe.to(device)
-    pipe.to(torch.bfloat16)
+    #**** BETTER WAY ****#
+    pipe.to(device, torch.bfloat16)
+    #**** BETTER WAY ****#
+
+    #pipe.to(device)
+    #pipe.to(torch.bfloat16)
 
     print(f'Pipeline: ')
     #print(f'_optional_components: {pipe._optional_components}')
@@ -276,6 +282,11 @@ def generate_30(
     num_inference_steps: int = 125,
     use_resolution_binning: bool = True, 
     lora_scale: float = 0.5,
+    b1:float = 1.0,
+    b2:float = 1.0,
+    s1:float = 1.0,
+    s2:float = 1.0,
+    freeu: bool = False,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     torch.set_default_device('cuda')
@@ -302,6 +313,9 @@ def generate_30(
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     batch_options = options.copy()
+    if freeu:
+        register_free_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
+        register_free_crossattn_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
     rv_image = pipe(**batch_options).images[0]
     sd_image_path = f"rv50_A_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
@@ -332,6 +346,11 @@ def generate_60(
     num_inference_steps: int = 250,
     use_resolution_binning: bool = True, 
     lora_scale: float = 0.5,
+    b1:float = 1.0,
+    b2:float = 1.0,
+    s1:float = 1.0,
+    s2:float = 1.0,
+    freeu: bool = False,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     #torch.set_default_device('cuda')
@@ -358,6 +377,9 @@ def generate_60(
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     batch_options = options.copy()
+    if freeu:
+        register_free_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
+        register_free_crossattn_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
     rv_image = pipe(**batch_options).images[0]
     sd_image_path = f"rv50_A_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
@@ -388,6 +410,11 @@ def generate_90(
     num_inference_steps: int = 250,
     use_resolution_binning: bool = True, 
     lora_scale: float = 0.5,
+    b1:float = 1.0,
+    b2:float = 1.0,
+    s1:float = 1.0,
+    s2:float = 1.0,
+    freeu: bool = False,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     #torch.set_default_device('cuda')
@@ -414,6 +441,9 @@ def generate_90(
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
     batch_options = options.copy()
+    if freeu:
+        register_free_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
+        register_free_crossattn_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
     rv_image = pipe(**batch_options).images[0]
     sd_image_path = f"rv50_A_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
@@ -537,6 +567,34 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                 step=10,
                 value=180,
             )
+        with gr.Row():
+                    freeu = gr.Checkbox(label="Use FreeU", value=False)
+
+                    b1 = gr.Slider(label='b1: backbone factor of the first stage block of decoder',
+                                            minimum=1,
+                                            maximum=1.6,
+                                            step=0.01,
+                                            value=1.1
+                                  )
+                    b2 = gr.Slider(label='b2: backbone factor of the second stage block of decoder',
+                                            minimum=1,
+                                            maximum=1.6,
+                                            step=0.01,
+                                            value=1.2
+                                  )
+        with gr.Row():
+                    s1 = gr.Slider(label='s1: skip factor of the first stage block of decoder',
+                                            minimum=0,
+                                            maximum=1,
+                                            step=0.1,
+                                            value=0.2
+                                  )
+                    s2 = gr.Slider(label='s2: skip factor of the second stage block of decoder',
+                                            minimum=0,
+                                            maximum=1,
+                                            step=0.1,
+                                            value=0.2
+                                  )    
 
     gr.Examples(
         examples=examples,
@@ -568,6 +626,11 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             lora_scale,
+            b1,
+            b2,
+            s1,
+            s2,
+            freeu,
         ],
         outputs=[result],
     )
@@ -589,6 +652,11 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             lora_scale,
+            b1,
+            b2,
+            s1,
+            s2,
+            freeu,
         ],
         outputs=[result],
     )
@@ -610,6 +678,11 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             lora_scale,
+            b1,
+            b2,
+            s1,
+            s2,
+            freeu,
         ],
         outputs=[result],
     )