Spaces:
Running on Zero
Running on Zero
| """Euler sampling with reference conditioning and classifier-free guidance. | |
| Vendored verbatim from https://github.com/Chenkin-x/ComfyUI-Moss-Image-Edit | |
| (moss_runtime/sampling.py, GPLv3). | |
| """ | |
| import math | |
| def sample_latents(editor, scheduler, reference_latents, positive, negative=None, *, | |
| height: int, width: int, steps: int = 50, cfg: float = 4, | |
| seed: int = 42, device="cuda", latent_channels: int = 128, | |
| use_kv_cache: bool = True, progress=None): | |
| """Euler sampling with one shared reference cache across steps and CFG text. | |
| Conditions expose ``embeds`` [1,S,D] and ``attention_mask`` [1,S]. Reference | |
| tensors are [1,C,h,w]. The wrapper receives normalized t; the scheduler | |
| receives its original 0..1000 timestep. All CFG/Euler arithmetic is FP32. | |
| """ | |
| import torch | |
| if min(height, width, steps, latent_channels) < 1 or height % 16 or width % 16: | |
| raise ValueError("Invalid latent dimensions/steps; image dimensions must be multiples of 16.") | |
| if not math.isfinite(cfg) or cfg < 0: | |
| raise ValueError("cfg must be finite and nonnegative.") | |
| if cfg != 1 and negative is None: | |
| raise ValueError("CFG requires negative text conditioning.") | |
| if not reference_latents: | |
| raise ValueError("Editing requires at least one reference image.") | |
| if getattr(scheduler.config, "num_train_timesteps", 1000) != 1000: | |
| raise ValueError("The edit wrapper requires the official 1000-timestep convention.") | |
| device = torch.device(device) | |
| references = [reference.detach().to(device) for reference in reference_latents] | |
| positive = positive.to(device) | |
| negative = None if negative is None else negative.to(device) | |
| generator = torch.Generator("cpu").manual_seed(seed) | |
| latents = torch.randn( | |
| (1, latent_channels, height // 16, width // 16), | |
| generator=generator, device="cpu", dtype=torch.float32, | |
| ).to(device) | |
| editor.eval() | |
| with torch.no_grad(): | |
| cache = editor.prepare_cache(references) if use_kv_cache else None | |
| scheduler.set_timesteps(steps, device=device) | |
| total = len(scheduler.timesteps) | |
| for index, timestep in enumerate(scheduler.timesteps): | |
| normalized = torch.as_tensor(timestep, device=device, dtype=torch.float32).reshape(1) / 1000 | |
| def predict(condition): | |
| return editor( | |
| hidden_states=latents, timestep=normalized, | |
| encoder_hidden_states=condition.embeds, | |
| reference_latents=references, | |
| encoder_attention_mask=condition.attention_mask, cache=cache, | |
| ).float() | |
| prediction = predict(positive) | |
| if cfg != 1: | |
| unconditional = predict(negative) | |
| prediction = unconditional + cfg * (prediction - unconditional) | |
| latents = scheduler.step(prediction, timestep, latents, return_dict=False)[0].float() | |
| if progress is not None: | |
| progress(index + 1, total) | |
| if not bool(torch.isfinite(latents).all()): | |
| raise RuntimeError("Sampling produced non-finite latents.") | |
| return latents | |