Update app.py

app.py

@@ -34,7 +34,7 @@ FTP_PASS = "GoogleBez12!"
 FTP_DIR = "1ink.us/stable_diff/"  # Remote directory on FTP server
 
 DESCRIPTIONXX = """
-    ## ⚡⚡⚡⚡ REALVISXL V5.0 BF16 (Tester B) ⚡⚡⚡⚡
+    ## ⚡⚡⚡⚡ REALVISXL V5.0 BF16 (Tester C) ⚡⚡⚡⚡
 """
 
 examples = [
@@ -78,7 +78,42 @@ STYLE_NAMES = list(styles.keys())
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    
+
+def scheduler_swap_callback(pipeline, step_index, timestep):
+        # adjust the batch_size of prompt_embeds according to guidance_scale
+    if step_index == int(pipeline.num_timesteps * 0.1):
+        print("-- swapping scheduler --")
+      #  pipeline.scheduler = euler_scheduler
+        torch.set_float32_matmul_precision("high")
+       # pipe.vae = vae_b
+        torch.backends.cudnn.allow_tf32 = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+        #  torch.backends.cudnn.deterministic = True
+        torch.backends.cuda.preferred_blas_library="cublaslt"
+    if step_index == int(pipeline.num_timesteps * 0.5):
+       # torch.set_float32_matmul_precision("medium")
+        pipe.unet.to(torch.float64)
+      #  pipe.guidance_scale=1.0
+       # pipe.scheduler.set_timesteps(num_inference_steps*.70)
+      #  print(f"-- setting step {pipeline.num_timesteps * 0.1} --")
+      #  pipeline.scheduler._step_index = pipeline.num_timesteps * 0.1 
+    if step_index == int(pipeline.num_timesteps * 0.9):
+        torch.backends.cuda.preferred_blas_library="cublas"
+        torch.backends.cudnn.allow_tf32 = False
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.set_float32_matmul_precision("highest")
+        pipe.unet.to(torch.bfloat16)
+     #   pipe.vae = vae_a
+     #   pipe.unet = unet_a
+      #  torch.backends.cudnn.deterministic = False
+        print("-- swapping scheduler --")
+      #  pipeline.scheduler = heun_scheduler
+        #pipe.scheduler.set_timesteps(num_inference_steps*.70)
+      #  print(f"-- setting step {pipeline.num_timesteps * 0.9} --")
+      #  pipeline.scheduler._step_index = pipeline.num_timesteps * 0.9
+    return
+
+
 def load_and_prepare_model():
     #vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False).to(device=device, dtype=torch.bfloat16)
     #vaeRV = AutoencoderKL.from_pretrained("SG161222/RealVisXL_V5.0", subfolder='vae', safety_checker=None, use_safetensors=False).to(device).to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
@@ -100,7 +135,7 @@ def load_and_prepare_model():
     return pipe
     
 # Preload and compile both models
-pipe =load_and_prepare_model()
+pipe = load_and_prepare_model()
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -172,6 +207,7 @@ def generate_30(
         "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
+        "callback_on_step_end": scheduler_swap_callback
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
@@ -212,6 +248,7 @@ def generate_60(
         "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
+        "callback_on_step_end": scheduler_swap_callback
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
@@ -252,6 +289,7 @@ def generate_90(
         "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
+        "callback_on_step_end": scheduler_swap_callback
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True