Update app.py

app.py

@@ -27,10 +27,7 @@ import time
 import datetime
 from gradio import themes
 from image_gen_aux import UpscaleWithModel
-
-#from diffusers.models.attention_processor import AttnProcessor2_0
-
-from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
+from diffusers.models.attention_processor import AttnProcessor2_0
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -171,6 +168,7 @@ 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****#
@@ -207,10 +205,11 @@ 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()
+    
     #**** BETTER WAY ****#
     pipe.to(device, torch.bfloat16)
     #**** BETTER WAY ****#
-
+    
     #pipe.to(device)
     #pipe.to(torch.bfloat16)
 
@@ -285,11 +284,6 @@ 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')
@@ -316,12 +310,6 @@ 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:
-        pipe.enable_freeu(s1,s2,b1,b2)
-        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)
-    else:
-        pipe.disable_freeu()
     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)
@@ -352,11 +340,6 @@ 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')
@@ -383,12 +366,6 @@ 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:
-        pipe.enable_freeu(s1,s2,b1,b2)
-        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)
-    else:
-        pipe.disable_freeu()
     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)
@@ -419,11 +396,6 @@ 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')
@@ -450,12 +422,6 @@ 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:
-        pipe.enable_freeu(s1,s2,b1,b2)
-        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)
-    else:
-        pipe.disable_freeu()
     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)
@@ -579,34 +545,6 @@ 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,
@@ -638,11 +576,6 @@ 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],
     )
@@ -664,11 +597,6 @@ 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],
     )
@@ -690,11 +618,6 @@ 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],
     )