Update app.py

app.py

@@ -316,7 +316,13 @@ def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         #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 {"latents": callback_kwargs["latents"]}   
+    if step_index == int(pipeline.num_timesteps * self.config.cutoff_step_ratio)):
+            prompt_embeds = callback_kwargs["prompt_embeds"]
+            prompt_embeds = prompt_embeds.chunk(2)[-1]
+                # update guidance_scale and prompt_embeds
+            pipeline._guidance_scale = new_scale
+            callback_kwargs["prompt_embeds"] = prompt_embeds
+    return callback_kwargs
 
 def upload_to_ftp(filename):
     try:
@@ -359,6 +365,7 @@ def generate_30(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    guidance_cutoff: float = 1.0,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
@@ -374,7 +381,8 @@ def generate_30(
         "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
-        "callback_on_step_end": pyx.scheduler_swap_callback
+        "callback_on_step_end": pyx.scheduler_swap_callback(cutoff_step_ratio=guidance_cutoff),
+        "callback_on_step_end_tensor_inputs": ['prompt_embeds']
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
@@ -412,6 +420,7 @@ def generate_60(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    guidance_cutoff: float = 1.0,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
@@ -427,7 +436,8 @@ def generate_60(
         "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
-        "callback_on_step_end": scheduler_swap_callback
+        "callback_on_step_end": pyx.scheduler_swap_callback(cutoff_step_ratio=guidance_cutoff),
+        "callback_on_step_end_tensor_inputs": ['prompt_embeds']
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
@@ -455,6 +465,7 @@ def generate_90(
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
+    guidance_cutoff: float = 1.0,
     use_resolution_binning: bool = True, 
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
@@ -470,7 +481,8 @@ def generate_90(
         "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
-        "callback_on_step_end": scheduler_swap_callback
+        "callback_on_step_end": pyx.scheduler_swap_callback(cutoff_step_ratio=guidance_cutoff),
+        "callback_on_step_end_tensor_inputs": ['prompt_embeds']
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
@@ -575,6 +587,13 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                 step=0.1,
                 value=3.8,
             )
+            guidance_cutoff = gr.Slider(
+                label="Guidance Scale Cutoff",
+                minimum=0.01,
+                maximum=1.0,
+                step=0.01,
+                value=1.0,
+            )
             num_inference_steps = gr.Slider(
                 label="Number of inference steps",
                 minimum=10,
@@ -610,6 +629,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             width,
             height,
             guidance_scale,
+            guidance_cutoff,
             num_inference_steps,
         ],
         outputs=[result],
@@ -629,6 +649,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             width,
             height,
             guidance_scale,
+            guidance_cutoff,
             num_inference_steps,
         ],
         outputs=[result],
@@ -648,6 +669,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             width,
             height,
             guidance_scale,
+            guidance_cutoff,
             num_inference_steps,
         ],
         outputs=[result],