This commit is contained in:
Haoming 2025-08-04 11:28:14 +08:00
parent 6dc921ecfc
commit 129faabc4c
4 changed files with 11 additions and 4 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)