cast to float

app.py

@@ -66,17 +66,16 @@ def gen(
         encode_prompts(sdxl_pipe, full_prompt, DEFAULT_NEGATIVE_PROMPT)
     )
     set_seed(seed)
-    with torch.autocast("cuda"):
-        result = sdxl_pipe(
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_embeds2,
-            negative_pooled_prompt_embeds=neg_pooled_embeds2,
-            num_inference_steps=24,
-            width=1024,
-            height=1024,
-            guidance_scale=6.0,
-        ).images[0]
+    result = sdxl_pipe(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        pooled_prompt_embeds=pooled_embeds2,
+        negative_pooled_prompt_embeds=neg_pooled_embeds2,
+        num_inference_steps=24,
+        width=1024,
+        height=1024,
+        guidance_scale=6.0,
+    ).images[0]
     torch.cuda.empty_cache()
     t1 = time_ns()
 
@@ -138,7 +137,9 @@ click "Next" button until you get the dragon girl you like.
                         value=list(DEFAULT_STYLE_LIST)[0],
                     )
                 submit = gr.Button("Next", variant="primary")
-                dtg_output = gr.TextArea(label="DTG output", lines=9, show_copy_button=True)
+                dtg_output = gr.TextArea(
+                    label="DTG output", lines=9, show_copy_button=True
+                )
                 cost_time = gr.Markdown()
             with gr.Column(scale=4):
                 result = gr.Image(label="Result", type="numpy", interactive=False)

diff.py

@@ -5,6 +5,8 @@ from diffusers import StableDiffusionXLKDiffusionPipeline
 from k_diffusion.sampling import get_sigmas_polyexponential
 from k_diffusion.sampling import sample_dpmpp_2m_sde
 
+torch.set_float32_matmul_precision("mediun")
+
 
 def set_timesteps_polyexponential(self, orig_sigmas, num_inference_steps, device=None):
     self.num_inference_steps = num_inference_steps
@@ -19,6 +21,16 @@ def set_timesteps_polyexponential(self, orig_sigmas, num_inference_steps, device
     self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas.new_zeros([1])])
 
 
+def model_forward(k_diffusion_model: torch.nn.Module):
+    orig_forward = k_diffusion_model.forward
+    def forward(*args, **kwargs):
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            result = orig_forward(*args, **kwargs)
+        return result.float()
+
+    return forward
+
+
 def load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon", device="cuda"):
     pipe: StableDiffusionXLKDiffusionPipeline
     pipe = StableDiffusionXLKDiffusionPipeline.from_pretrained(
@@ -28,6 +40,7 @@ def load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon", device="cuda"):
         set_timesteps_polyexponential, pipe.scheduler, pipe.scheduler.sigmas
     )
     pipe.sampler = partial(sample_dpmpp_2m_sde, eta=0.35, solver_type="heun")
+    pipe.k_diffusion_model.forward = model_forward(pipe.k_diffusion_model)
     return pipe
 
 
@@ -104,4 +117,3 @@ def encode_prompts(pipe: StableDiffusionXLKDiffusionPipeline, prompt, neg_prompt
     neg_pooled_embeds2 = torch.mean(torch.stack(neg_pooled_embeds2, dim=0), dim=0)
 
     return prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2
-