From 129faabc4ccc8785de9e8a257bc3338eafac3003 Mon Sep 17 00:00:00 2001 From: Haoming Date: Mon, 4 Aug 2025 11:28:14 +0800 Subject: [PATCH] kontext --- backend/loader.py | 3 +++ backend/sampling/sampling_function.py | 8 ++++++-- modules/sd_samplers_kdiffusion.py | 2 +- modules/sd_samplers_timesteps.py | 2 +- 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/backend/loader.py b/backend/loader.py index f4b0a016..8b788843 100644 --- a/backend/loader.py +++ b/backend/loader.py @@ -9,6 +9,7 @@ from transformers import modeling_utils import backend.args from backend import memory_management +from backend.args import dynamic_args from backend.diffusion_engine.chroma import Chroma from backend.diffusion_engine.flux import Flux from backend.diffusion_engine.sd15 import StableDiffusion @@ -535,6 +536,8 @@ def forge_loader(sd: os.PathLike, additional_state_dicts: list[os.PathLike] = No else: huggingface_components["scheduler"].config.prediction_type = prediction_types.get(estimated_config.model_type.name, huggingface_components["scheduler"].config.prediction_type) + dynamic_args["kontext"] = "kontext" in str(sd).lower() + for M in possible_models: if any(isinstance(estimated_config, x) for x in M.matched_guesses): return M(estimated_config=estimated_config, huggingface_components=huggingface_components) diff --git a/backend/sampling/sampling_function.py b/backend/sampling/sampling_function.py index f207635f..0cd926b1 100644 --- a/backend/sampling/sampling_function.py +++ b/backend/sampling/sampling_function.py @@ -367,8 +367,12 @@ def sampling_function(self, denoiser_params, cond_scale, cond_composition): return denoised, cond_pred, uncond_pred -def sampling_prepare(unet: "UnetPatcher", x: torch.Tensor): - unet.set_transformer_option("ref_latents", [x.detach().clone()]) +def sampling_prepare(unet: "UnetPatcher", x: torch.Tensor, *, is_img2img: bool = False): + if is_img2img and dynamic_args.get("kontext", False): + unet.set_transformer_option("ref_latents", [x.detach().clone()]) + else: + unet.set_transformer_option("ref_latents", None) + B, C, H, W = x.shape memory_estimation_function = unet.model_options.get("memory_peak_estimation_modifier", unet.memory_required) diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 0e4072e2..4b736bff 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -135,7 +135,7 @@ class KDiffusionSampler(sd_samplers_common.Sampler): def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None): unet_patcher = self.model_wrap.inner_model.forge_objects.unet - sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x) + sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x, is_img2img=True) steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps) diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py index 514e1d63..92c09dff 100644 --- a/modules/sd_samplers_timesteps.py +++ b/modules/sd_samplers_timesteps.py @@ -73,7 +73,7 @@ class CompVisSampler(sd_samplers_common.Sampler): def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None): unet_patcher = self.model_wrap.inner_model.forge_objects.unet - sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x) + sampling_prepare(self.model_wrap.inner_model.forge_objects.unet, x=x, is_img2img=True) self.model_wrap.inner_model.alphas_cumprod = self.model_wrap.inner_model.alphas_cumprod.to(x.device)