Fix refiner progress

README.md

@@ -48,3 +48,37 @@ preload_from_hub:
 # diffusion-xl
 
 Gradio app for Stable Diffusion XL.
+
+## Usage
+
+See [DOCS.md](https://huggingface.co/spaces/adamelliotfields/diffusion-xl/blob/main/DOCS.md).
+
+## Installation
+
+```sh
+# clone
+git clone https://huggingface.co/spaces/adamelliotfields/diffusion-xl.git
+cd diffusion-xl
+git remote set-url origin https://adamelliotfields:$HF_TOKEN@huggingface.co/spaces/adamelliotfields/diffusion-xl
+
+# install
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+
+# gradio
+python app.py --port 7860
+```
+
+## Development
+
+See [pull requests and discussions](https://huggingface.co/docs/hub/en/repositories-pull-requests-discussions).
+
+```sh
+git fetch origin refs/pr/42:pr/42
+git checkout pr/42
+# ...
+git add .
+git commit -m "Commit message"
+git push origin pr/42:refs/pr/42
+```

lib/config.py

@@ -6,7 +6,8 @@ from diffusers import (
     DPMSolverMultistepScheduler,
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
-    PNDMScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    KDPM2DiscreteScheduler,
     StableDiffusionXLImg2ImgPipeline,
     StableDiffusionXLPipeline,
 )
@@ -34,22 +35,23 @@ Config = SimpleNamespace(
     ],
     VAE_MODEL="madebyollin/sdxl-vae-fp16-fix",
     REFINER_MODEL="stabilityai/stable-diffusion-xl-refiner-1.0",
-    SCHEDULER="DEIS 2M",
+    SCHEDULER="Euler",
     SCHEDULERS={
         "DDIM": DDIMScheduler,
         "DEIS 2M": DEISMultistepScheduler,
         "DPM++ 2M": DPMSolverMultistepScheduler,
+        "DPM2": KDPM2DiscreteScheduler,
+        "DPM2 a": KDPM2AncestralDiscreteScheduler,
         "Euler": EulerDiscreteScheduler,
         "Euler a": EulerAncestralDiscreteScheduler,
-        "PNDM": PNDMScheduler,
     },
     STYLE="sai-enhance",
     WIDTH=896,
     HEIGHT=1152,
     NUM_IMAGES=1,
     SEED=-1,
-    GUIDANCE_SCALE=5,
-    INFERENCE_STEPS=40,
+    GUIDANCE_SCALE=6,
+    INFERENCE_STEPS=35,
     DEEPCACHE_INTERVAL=1,
     SCALE=1,
     SCALES=[1, 2, 4],

lib/inference.py

@@ -91,7 +91,7 @@ def generate(
     style=None,
     seed=None,
     model="stabilityai/stable-diffusion-xl-base-1.0",
-    scheduler="DEIS 2M",
+    scheduler="DDIM",
     width=1024,
     height=1024,
     guidance_scale=7.5,
@@ -100,7 +100,7 @@ def generate(
     scale=1,
     num_images=1,
     use_karras=False,
-    use_refiner=True,
+    use_refiner=False,
     Info: Callable[[str], None] = None,
     Error=Exception,
     progress=None,
@@ -112,30 +112,39 @@ def generate(
     if seed is None or seed < 0:
         seed = int(datetime.now().timestamp() * 1_000_000) % (2**64)
 
-    EMBEDDINGS_TYPE = ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED
-
     KIND = "txt2img"
-
-    CURRENT_IMAGE = 1
     CURRENT_STEP = 0
+    CURRENT_IMAGE = 1
+    EMBEDDINGS_TYPE = ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED
 
     if progress is not None:
         TQDM = False
-        progress((0, inference_steps), desc=f"Generating image {CURRENT_IMAGE}/{num_images}")
+        progress((0, inference_steps), desc=f"Generating image 1/{num_images}")
     else:
         TQDM = True
 
     def callback_on_step_end(pipeline, step, timestep, latents):
         nonlocal CURRENT_IMAGE, CURRENT_STEP
+
         if progress is None:
             return latents
+
         strength = 1
         total_steps = min(int(inference_steps * strength), inference_steps)
-        CURRENT_STEP += step + 1
+
+        # if steps are different we're in the refiner
+        refining = False
+        if CURRENT_STEP == step:
+            CURRENT_STEP = step + 1
+        else:
+            refining = True
+            CURRENT_STEP += 1
+
         progress(
             (CURRENT_STEP, total_steps),
-            desc=f"Generating image {CURRENT_IMAGE}/{num_images}",
+            desc=f"{'Refining' if refining else 'Generating'} image {CURRENT_IMAGE}/{num_images}",
         )
+
         return latents
 
     start = time.perf_counter()
@@ -150,6 +159,7 @@ def generate(
         use_refiner,
         TQDM,
     )
+
     # prompt embeds for base and refiner
     compel_1 = Compel(
         text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
@@ -232,7 +242,6 @@ def generate(
 
             if progress is not None:
                 refiner_kwargs["callback_on_step_end"] = callback_on_step_end
-
             if use_refiner:
                 image = refiner(**refiner_kwargs).images[0]
             if scale > 1:

lib/loader.py

@@ -53,13 +53,13 @@ class Loader:
 
     def _unload(self, model, refiner, scale):
         to_unload = []
-        if self._should_unload_upscaler(scale):
-            to_unload.append("upscaler")
-        if self._should_unload_refiner(refiner):
-            to_unload.append("refiner")
         if self._should_unload_pipeline(model):
             to_unload.append("model")
             to_unload.append("pipe")
+        if self._should_unload_refiner(refiner):
+            to_unload.append("refiner")
+        if self._should_unload_upscaler(scale):
+            to_unload.append("upscaler")
         for component in to_unload:
             delattr(self, component)
         self._flush()
@@ -71,8 +71,13 @@ class Loader:
         if self.pipe is None:
             try:
                 print(f"Loading {model}...")
-                self.pipe = pipeline.from_pretrained(model, **kwargs).to("cuda")
                 self.model = model
+                self.pipe = pipeline.from_pretrained(model, **kwargs).to("cuda")
+                if self.refiner is not None:
+                    self.refiner.vae = self.pipe.vae
+                    self.refiner.scheduler = self.pipe.scheduler
+                    self.refiner.tokenizer_2 = self.pipe.tokenizer_2
+                    self.refiner.text_encoder_2 = self.pipe.text_encoder_2
             except Exception as e:
                 print(f"Error loading {model}: {e}")
                 self.model = None
@@ -125,7 +130,7 @@ class Loader:
             "steps_offset": 1,
         }
 
-        if scheduler not in ["DDIM", "Euler a", "PNDM"]:
+        if scheduler not in ["DDIM", "Euler a"]:
             scheduler_kwargs["use_karras_sigmas"] = karras
 
         # https://github.com/huggingface/diffusers/blob/8a3f0c1/scripts/convert_original_stable_diffusion_to_diffusers.py#L939
@@ -169,6 +174,7 @@ class Loader:
                 print(f"{'Enabling' if karras else 'Disabling'} Karras sigmas...")
             if not same_scheduler or not same_karras:
                 self.pipe.scheduler = Config.SCHEDULERS[scheduler](**scheduler_kwargs)
+                self.refiner.scheduler = self.pipe.scheduler
 
         # https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/model_index.json
         refiner_kwargs = {