| 4 | |
| 5 | |
| 6 | def get_sample_forward(attn_bank, save_steps, single_layers, double_layers): |
| 7 | @torch.no_grad() |
| 8 | def sample_forward(model, x, sigmas, extra_args=None, callback=None, disable=None): |
| 9 | attn_bank.clear() |
| 10 | attn_bank["save_steps"] = save_steps |
| 11 | |
| 12 | extra_args = {} if extra_args is None else extra_args |
| 13 | |
| 14 | model_options = extra_args.get("model_options", {}) |
| 15 | model_options = {**model_options} |
| 16 | transformer_options = model_options.get("transformer_options", {}) |
| 17 | transformer_options = {**transformer_options} |
| 18 | model_options["transformer_options"] = transformer_options |
| 19 | extra_args["model_options"] = model_options |
| 20 | |
| 21 | N = len(sigmas) - 1 |
| 22 | s_in = x.new_ones([x.shape[0]]) |
| 23 | for i in trange(N, disable=disable): |
| 24 | sigma = sigmas[i] |
| 25 | sigma_next = sigmas[i + 1] |
| 26 | |
| 27 | if N - i - 1 < save_steps: |
| 28 | attn_bank[N - i - 1] = {"first": {}, "mid": {}} |
| 29 | |
| 30 | transformer_options["rfedit"] = { |
| 31 | "step": N - i - 1, |
| 32 | "process": "forward" if N - i - 1 < save_steps else None, |
| 33 | "pred": "first", |
| 34 | "bank": attn_bank, |
| 35 | "single_layers": single_layers, |
| 36 | "double_layers": double_layers, |
| 37 | } |
| 38 | |
| 39 | pred = model(x, s_in * sigma, **extra_args) |
| 40 | |
| 41 | transformer_options["rfedit"] = { |
| 42 | "step": N - i - 1, |
| 43 | "process": "forward" if N - i - 1 < save_steps else None, |
| 44 | "pred": "mid", |
| 45 | "bank": attn_bank, |
| 46 | "single_layers": single_layers, |
| 47 | "double_layers": double_layers, |
| 48 | } |
| 49 | |
| 50 | img_mid = x + (sigma_next - sigma) / 2 * pred |
| 51 | sigma_mid = sigma + (sigma_next - sigma) / 2 |
| 52 | pred_mid = model(img_mid, s_in * sigma_mid, **extra_args) |
| 53 | |
| 54 | first_order = (pred_mid - pred) / ((sigma_next - sigma) / 2) |
| 55 | x = ( |
| 56 | x |
| 57 | + (sigma_next - sigma) * pred |
| 58 | + 0.5 * (sigma_next - sigma) ** 2 * first_order |
| 59 | ) |
| 60 | |
| 61 | if callback is not None: |
| 62 | callback( |
| 63 | { |