Update app.py

app.py

@@ -5,10 +5,13 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline, EDMEulerScheduler, AutoencoderKL, DPMSolverMultistepScheduler
-from huggingface_hub import hf_hub_download, InferenceClient
+from diffusers import StableDiffusionXLImg2ImgPipeline, EDMEulerScheduler, AutoencoderKL
+from huggingface_hub import hf_hub_download
 
+# Load the VAE
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+
+# Download and load the model
 pipe_edit = StableDiffusionXLImg2ImgPipeline.from_single_file(
     hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors"),
     num_in_channels=8,
@@ -16,21 +19,28 @@ pipe_edit = StableDiffusionXLImg2ImgPipeline.from_single_file(
     vae=vae,
     torch_dtype=torch.float16,
 )
+
+# Set the scheduler
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_edit.to("cuda")
 
-refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+# Load the refiner
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0", 
+    vae=vae, 
+    torch_dtype=torch.float16, 
+    use_safetensors=True, 
+    variant="fp16"
+)
 refiner.to("cuda")
 
+# Patch for the scheduler
 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)
+    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
@@ -38,16 +48,8 @@ def set_timesteps_patched(self, num_inference_steps: int, device=None):
 
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
 
-def king(
-        input_image,
-        instruction: str,
-        negative_prompt: str = "",
-        steps: int = 25,
-        randomize_seed: bool = True,
-        seed: int = 2404,
-        guidance_scale: float = 6,
-        progress=gr.Progress(track_tqdm=True)
-    ):
+# Function to perform image editing
+def king(input_image, instruction: str, negative_prompt: str = "", steps: int = 25, randomize_seed: bool = True, seed: int = 2404, guidance_scale: float = 6, progress=gr.Progress(track_tqdm=True)):
     input_image = Image.open(input_image).convert('RGB')
     if randomize_seed:
         seed = random.randint(0, 999999)
@@ -74,6 +76,7 @@ def king(
     ).images[0]  
     return seed, refine
 
+# CSS for the Gradio interface
 css = '''
 .gradio-container{max-width: 700px !important}
 h1{text-align:center}
@@ -82,17 +85,13 @@ footer {
 }
 '''
 
+# Examples for the Gradio interface
 examples = [
-    [
-        "./supercar.png",
-        "make it red",
-    ],
-    [
-        "./red_car.png",
-        "add some snow",
-    ],
+    ["./supercar.png", "make it red"],
+    ["./red_car.png", "add some snow"],
 ]
 
+# Creating the Gradio interface
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# Image Editing\n### Note: First image generation takes time")
     with gr.Row():
@@ -131,4 +130,4 @@ with gr.Blocks(css=css) as demo:
         outputs=[seed, input_image],
     )
 
-demo.queue(max_size=500).launch()
+demo.queue(max_size=500).launch()