Update app.py

app.py

@@ -11,20 +11,6 @@ from PIL import Image
 edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
 normal_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl.safetensors")
 
-def set_timesteps_patched(self, num_inference_steps: int, device = None):
-    self.num_inference_steps = num_inference_steps
-    
-    ramp = np.linspace(0, 1, self.num_inference_steps)
-    sigmas = torch.linspace(math.log(self.config.sigma_min), math.log(self.config.sigma_max), len(ramp)).exp().flip(0)
-    
-    sigmas = (sigmas).to(dtype=torch.float32, device=device)
-    self.timesteps = self.precondition_noise(sigmas)
-    
-    self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
-    self._step_index = None
-    self._begin_index = None
-    self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
-
 def resize_image(image, resolution):
     original_width, original_height = image.size
     
@@ -38,18 +24,17 @@ def resize_image(image, resolution):
     resized_img = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
     return resized_img
 
-EDMEulerScheduler.set_timesteps = set_timesteps_patched
 
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
     edit_file, num_in_channels=8, is_cosxl_edit=True, vae=vae, torch_dtype=torch.float16,
 )
-pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
+pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction", sigma_schedule="exponential")
 pipe_edit.to("cuda")
 
 pipe_normal = StableDiffusionXLPipeline.from_single_file(normal_file, torch_dtype=torch.float16, vae=vae)
-pipe_normal.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
+pipe_normal.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction", sigma_schedule="exponential"s)
 pipe_normal.to("cuda")
 
 @spaces.GPU