1{
2  "op": "replace",
3  "filename": "CFGpp-main/latent_diffusion.py",
4  "start_line": 621,
5  "end_line": 679,
6  "content": "@register_solver(\"ddim_cfg++\")\nclass BaseDDIMCFGpp(StableDiffusion):\n    \"\"\"\n    DPM++ 2M with Adaptive Renoising Interpolation (ARI).\n    Combines second-order multistep ODE solving with progressive\n    interpolation between CFG++ and standard CFG renoising.\n    Early steps use CFG++ (unconditional noise) for trajectory stability,\n    later steps blend in conditional guidance for improved text alignment.\n    \"\"\"\n    def __init__(self, solver_config: Dict, model_key: str = \"runwayml/stable-diffusion-v1-5\",\n                 device: Optional[torch.device] = None, **kwargs):\n        super().__init__(solver_config, model_key, device, **kwargs)\n        self.ari_max_blend = 0.25  # Maximum blend toward conditional denoised\n\n    @torch.autocast(device_type='cuda', dtype=torch.float16)\n    def sample(self, cfg_guidance=7.5, prompt=[\"\", \"\"], callback_fn=None, **kwargs):\n        t_fn = lambda sigma: sigma.log().neg()\n\n        uc, c = self.get_text_embed(null_prompt=prompt[0], prompt=prompt[1])\n        total_sigmas = (1 - self.total_alphas).sqrt() / self.total_alphas.sqrt()\n        sigmas = get_sigmas_karras(len(self.scheduler.timesteps), total_sigmas.min(), total_sigmas.max(), rho=7.)\n\n        x = self.initialize_latent(method=\"random_kdiffusion\",\n                                   latent_dim=(1, 4, 64, 64), sigmas=sigmas).to(torch.float16)\n        old_blend = None\n        num_steps = len(self.scheduler.timesteps)\n\n        pbar = tqdm(self.scheduler.timesteps, desc=\"DPM++2M-ARI\")\n        for i, _ in enumerate(pbar):\n            sigma = sigmas[i]\n            new_t = self.timestep(sigma).to(self.device)\n\n            with torch.no_grad():\n                denoised, uncond_denoised = self.kdiffusion_x_to_denoised(\n                    x, sigma, uc, c, cfg_guidance, new_t)\n\n            # Adaptive Renoising Interpolation (ARI):\n            # Cosine schedule blending from pure CFG++ toward partial standard CFG\n            progress = i / max(num_steps - 1, 1)\n            w = 0.5 * (1.0 - (progress * 3.14159265)).__class__((progress * 3.14159265).__neg__()).__class__(1.0)\n            # Use simple quadratic ramp instead for numerical stability\n            w = (progress ** 2) * self.ari_max_blend\n            blend = (1.0 - w) * uncond_denoised + w * denoised\n\n            # DPM++ 2M multistep update with ARI-blended renoising direction\n            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])\n            h = t_next - t\n            if old_blend is None or sigmas[i + 1] == 0:\n                x = denoised + self.to_d(x, sigmas[i], blend) * sigmas[i + 1]\n            else:\n                h_last = t - t_fn(sigmas[i - 1])\n                r = h_last / h\n                extra1 = -torch.exp(-h) * blend - (-h).expm1() * (blend - old_blend) / (2 * r)\n                extra2 = torch.exp(-h) * x\n                x = denoised + extra1 + extra2\n            old_blend = blend\n\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': denoised.detach(), 'zt': x.detach(), 'decode': self.decode}\n                callback_kwargs = callback_fn(i, new_t, callback_kwargs)\n                denoised = callback_kwargs[\"z0t\"]\n                x = callback_kwargs[\"zt\"]\n\n        img = self.decode(x)\n        img = (img / 2 + 0.5).clamp(0, 1)\n        return img.detach().cpu()\n\n\n"
7}

1{
2  "content": "@register_solver(\"ddim_cfg++\")\nclass BaseDDIMCFGpp(SDXL):\n    \"\"\"DPM++ 2M with Adaptive Renoising Interpolation (ARI) for SDXL.\"\"\"\n    quantize = True\n\n    def reverse_process(self, null_prompt_embeds, prompt_embeds, cfg_guidance,\n                        add_cond_kwargs, shape=(1024, 1024), callback_fn=None, **kwargs):\n        t_fn = lambda sigma: sigma.log().neg()\n        ari_max_blend = 0.25\n\n        alphas = self.scheduler.alphas_cumprod[self.scheduler.timesteps.int().cpu()].cpu()\n        sigmas = (1 - alphas).sqrt() / alphas.sqrt()\n\n        x = self.initialize_latent(size=(1, 4, shape[1] // self.vae_scale_factor,\n                                         shape[0] // self.vae_scale_factor))\n        x = x * sigmas[0]\n\n        old_blend = None\n        num_steps = len(self.scheduler.timesteps) - 1\n\n        pbar = tqdm(self.scheduler.timesteps[:-1].int(), desc='SDXL-DPM++2M-ARI')\n        for i, _ in enumerate(pbar):\n            at = alphas[i]\n            sigma = sigmas[i]\n            c_in, c_out = at.sqrt(), -sigma\n\n            new_t = self.sigma_to_t(sigma).to(self.device)\n            with torch.no_grad():\n                noise_uc, noise_c = self.predict_noise(x * c_in, new_t, null_prompt_embeds,\n                                                        prompt_embeds, add_cond_kwargs)\n                noise_pred = noise_uc + cfg_guidance * (noise_c - noise_uc)\n\n            denoised = x + c_out * noise_pred\n            uncond_denoised = x + c_out * noise_uc\n\n            # Adaptive Renoising Interpolation\n            progress = i / max(num_steps - 1, 1)\n            w = (progress ** 2) * ari_max_blend\n            blend = (1.0 - w) * uncond_denoised + w * denoised\n\n            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])\n            h = t_next - t\n            if old_blend is None or sigmas[i + 1] == 0:\n                x = denoised + self.to_d(x, sigmas[i], blend) * sigmas[i + 1]\n            else:\n                h_last = t - t_fn(sigmas[i - 1])\n                r = h_last / h\n                extra1 = -torch.exp(-h) * blend - (-h).expm1() * (blend - old_blend) / (2 * r)\n                extra2 = torch.exp(-h) * x\n                x = denoised + extra1 + extra2\n            old_blend = blend\n\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': denoised.detach(), 'zt': x.detach(), 'decode': self.decode}\n                callback_kwargs = callback_fn(i, new_t, callback_kwargs)\n        return x\n\n\n",
3  "end_line": 755,
4  "filename": "CFGpp-main/latent_sdxl.py",
5  "op": "replace",
6  "start_line": 713
7}

1{
2  "content": "@register_solver(\"ddim_cfg++\")\nclass BaseDDIMCFGpp(StableDiffusion):\n    \"\"\"\n    DPM++ 2M with Adaptive Renoising Interpolation (ARI).\n    Combines second-order multistep ODE solving with progressive\n    interpolation between CFG++ and standard CFG renoising.\n    For low guidance (cfg~1), uses mostly conditional denoised for renoising.\n    For high guidance, starts with unconditional and progressively blends conditional.\n    \"\"\"\n    def __init__(self, solver_config: Dict, model_key: str = \"runwayml/stable-diffusion-v1-5\",\n                 device: Optional[torch.device] = None, **kwargs):\n        super().__init__(solver_config, model_key, device, **kwargs)\n\n    @torch.autocast(device_type='cuda', dtype=torch.float16)\n    def sample(self, cfg_guidance=7.5, prompt=[\"\", \"\"], callback_fn=None, **kwargs):\n        t_fn = lambda sigma: sigma.log().neg()\n\n        uc, c = self.get_text_embed(null_prompt=prompt[0], prompt=prompt[1])\n        total_sigmas = (1 - self.total_alphas).sqrt() / self.total_alphas.sqrt()\n        sigmas = get_sigmas_karras(len(self.scheduler.timesteps), total_sigmas.min(), total_sigmas.max(), rho=7.)\n\n        x = self.initialize_latent(method=\"random_kdiffusion\",\n                                   latent_dim=(1, 4, 64, 64), sigmas=sigmas).to(torch.float16)\n        old_blend = None\n        num_steps = len(self.scheduler.timesteps)\n\n        # Adaptive base blend: high cfg -> use more uncond (CFG++), low cfg -> use cond\n        # When cfg=1, the uncond noise direction is less informative, prefer cond denoised\n        base_blend = min(1.0, max(0.3, 1.0 - 0.7 / max(cfg_guidance, 0.1)))\n        # base_blend: fraction of conditional denoised in the renoising blend\n        # cfg=1 -> base_blend=0.3, cfg=7.5 -> base_blend~0.91\n\n        pbar = tqdm(self.scheduler.timesteps, desc=\"DPM++2M-ARI\")\n        for i, _ in enumerate(pbar):\n            sigma = sigmas[i]\n            new_t = self.timestep(sigma).to(self.device)\n\n            with torch.no_grad():\n                denoised, uncond_denoised = self.kdiffusion_x_to_denoised(\n                    x, sigma, uc, c, cfg_guidance, new_t)\n\n            # Adaptive Renoising Interpolation (ARI):\n            # Progressively increase conditional blend from early to late steps\n            progress = i / max(num_steps - 1, 1)\n            # Start with more unconditional (stability), end with more conditional (alignment)\n            w = base_blend * (0.5 + 0.5 * progress)\n            blend = (1.0 - w) * uncond_denoised + w * denoised\n\n            # DPM++ 2M multistep update with ARI-blended renoising direction\n            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])\n            h = t_next - t\n            if old_blend is None or sigmas[i + 1] == 0:\n                x = denoised + self.to_d(x, sigmas[i], blend) * sigmas[i + 1]\n            else:\n                h_last = t - t_fn(sigmas[i - 1])\n                r = h_last / h\n                extra1 = -torch.exp(-h) * blend - (-h).expm1() * (blend - old_blend) / (2 * r)\n                extra2 = torch.exp(-h) * x\n                x = denoised + extra1 + extra2\n            old_blend = blend\n\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': denoised.detach(), 'zt': x.detach(), 'decode': self.decode}\n                callback_kwargs = callback_fn(i, new_t, callback_kwargs)\n                denoised = callback_kwargs[\"z0t\"]\n                x = callback_kwargs[\"zt\"]\n\n        img = self.decode(x)\n        img = (img / 2 + 0.5).clamp(0, 1)\n        return img.detach().cpu()\n\n\n",
3  "end_line": 688,
4  "filename": "CFGpp-main/latent_diffusion.py",
5  "op": "replace",
6  "start_line": 621
7}

1{
2  "content": "@register_solver(\"ddim_cfg++\")\nclass BaseDDIMCFGpp(SDXL):\n    \"\"\"DPM++ 2M with Adaptive Renoising Interpolation (ARI) for SDXL.\"\"\"\n    quantize = True\n\n    def reverse_process(self, null_prompt_embeds, prompt_embeds, cfg_guidance,\n                        add_cond_kwargs, shape=(1024, 1024), callback_fn=None, **kwargs):\n        t_fn = lambda sigma: sigma.log().neg()\n\n        alphas = self.scheduler.alphas_cumprod[self.scheduler.timesteps.int().cpu()].cpu()\n        sigmas = (1 - alphas).sqrt() / alphas.sqrt()\n\n        x = self.initialize_latent(size=(1, 4, shape[1] // self.vae_scale_factor,\n                                         shape[0] // self.vae_scale_factor))\n        x = x * sigmas[0]\n\n        old_blend = None\n        num_steps = len(self.scheduler.timesteps) - 1\n        # Adaptive base blend: high cfg -> more uncond (CFG++), low cfg -> more cond\n        base_blend = min(1.0, max(0.3, 1.0 - 0.7 / max(cfg_guidance, 0.1)))\n\n        pbar = tqdm(self.scheduler.timesteps[:-1].int(), desc='SDXL-DPM++2M-ARI')\n        for i, _ in enumerate(pbar):\n            at = alphas[i]\n            sigma = sigmas[i]\n            c_in, c_out = at.sqrt(), -sigma\n\n            new_t = self.sigma_to_t(sigma).to(self.device)\n            with torch.no_grad():\n                noise_uc, noise_c = self.predict_noise(x * c_in, new_t, null_prompt_embeds,\n                                                        prompt_embeds, add_cond_kwargs)\n                noise_pred = noise_uc + cfg_guidance * (noise_c - noise_uc)\n\n            denoised = x + c_out * noise_pred\n            uncond_denoised = x + c_out * noise_uc\n\n            # Adaptive Renoising Interpolation\n            progress = i / max(num_steps - 1, 1)\n            w = base_blend * (0.5 + 0.5 * progress)\n            blend = (1.0 - w) * uncond_denoised + w * denoised\n\n            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])\n            h = t_next - t\n            if old_blend is None or sigmas[i + 1] == 0:\n                x = denoised + self.to_d(x, sigmas[i], blend) * sigmas[i + 1]\n            else:\n                h_last = t - t_fn(sigmas[i - 1])\n                r = h_last / h\n                extra1 = -torch.exp(-h) * blend - (-h).expm1() * (blend - old_blend) / (2 * r)\n                extra2 = torch.exp(-h) * x\n                x = denoised + extra1 + extra2\n            old_blend = blend\n\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': denoised.detach(), 'zt': x.detach(), 'decode': self.decode}\n                callback_kwargs = callback_fn(i, new_t, callback_kwargs)\n        return x\n\n\n",
3  "end_line": 770,
4  "filename": "CFGpp-main/latent_sdxl.py",
5  "op": "replace",
6  "start_line": 713
7}

1{
2  "op": "replace",
3  "filename": "CFGpp-main/latent_diffusion.py",
4  "start_line": 621,
5  "end_line": 679,
6  "content": "@register_solver(\"ddim_cfg++\")\nclass BaseDDIMCFGpp(StableDiffusion):\n    \"\"\"\n    CFG++ with Rescaled Guidance, Momentum Correction, and Dynamic Thresholding.\n    \n    Novel contributions over standard DDIM/CFG++:\n    1. Rescaled CFG: normalize guided noise to match conditional std, preventing\n       oversaturation that degrades CLIP alignment.\n    2. CFG++ renoising: use unconditional noise for the diffusion step (stabler ODE).\n    3. Momentum-corrected z0t: EMA of clean-image estimates across steps to reduce\n       high-frequency prediction jitter, especially in early noisy steps.\n    4. Dynamic thresholding (Imagen-style): clamp predicted clean latents to\n       prevent extreme values that cause artifacts.\n    \"\"\"\n    def __init__(self, solver_config, model_key=\"runwayml/stable-diffusion-v1-5\",\n                 device=None, **kwargs):\n        super().__init__(solver_config, model_key, device, **kwargs)\n        self.rescale_phi = 0.7        # CFG rescale blending factor\n        self.momentum_beta = 0.12     # EMA momentum for z0t correction\n        self.thresh_pct = 0.995       # dynamic thresholding percentile\n\n    @torch.autocast(device_type='cuda', dtype=torch.float16)\n    def sample(self, cfg_guidance=7.5, prompt=[\"\",\"\"], callback_fn=None, **kwargs):\n        uc, c = self.get_text_embed(null_prompt=prompt[0], prompt=prompt[1])\n        zt = self.initialize_latent().requires_grad_()\n        z0t_prev = None\n\n        pbar = tqdm(self.scheduler.timesteps, desc=\"CFG++R\")\n        for step, t in enumerate(pbar):\n            at = self.alpha(t)\n            at_prev = self.alpha(t - self.skip)\n            with torch.no_grad():\n                noise_uc, noise_c = self.predict_noise(zt, t, uc, c)\n                noise_cfg = noise_uc + cfg_guidance * (noise_c - noise_uc)\n                # Rescaled CFG: match std of conditional prediction\n                std_cfg = noise_cfg.std(dim=[1,2,3], keepdim=True) + 1e-8\n                std_c = noise_c.std(dim=[1,2,3], keepdim=True)\n                noise_rescaled = noise_cfg * (std_c / std_cfg)\n                noise_pred = self.rescale_phi * noise_rescaled + (1 - self.rescale_phi) * noise_cfg\n            # Tweedie: predict clean image\n            z0t = (zt - (1-at).sqrt() * noise_pred) / at.sqrt()\n            # Dynamic thresholding\n            s = torch.quantile(z0t.abs().reshape(z0t.shape[0], -1), self.thresh_pct, dim=1)\n            s = torch.clamp(s, min=1.0)[:, None, None, None]\n            z0t = z0t.clamp(-s, s) / s\n            # Momentum correction: EMA of z0t predictions (decay with progress)\n            if z0t_prev is not None:\n                progress = step / len(self.scheduler.timesteps)\n                beta = self.momentum_beta * (1.0 - progress)\n                z0t = (1 - beta) * z0t + beta * z0t_prev\n            z0t_prev = z0t.detach().clone()\n            # CFG++ renoising: use unconditional noise for stable ODE trajectory\n            zt = at_prev.sqrt() * z0t + (1-at_prev).sqrt() * noise_uc\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': z0t.detach(), 'zt': zt.detach(),\n                                    'decode': self.decode}\n                callback_kwargs = callback_fn(step, t, callback_kwargs)\n                z0t = callback_kwargs[\"z0t\"]\n                zt = callback_kwargs[\"zt\"]\n\n"
7}

1{
2  "op": "replace",
3  "filename": "CFGpp-main/latent_diffusion.py",
4  "start_line": 624,
5  "end_line": 682,
6  "content": "    #\n    # CFG++ with Rescaled Guidance, Momentum Correction & Dynamic Thresholding.\n    #\n    # Novel contributions over standard DDIM/CFG++:\n    # 1. Rescaled CFG: normalize guided noise to match conditional std,\n    #    preventing oversaturation that degrades CLIP alignment.\n    # 2. CFG++ renoising: use unconditional noise for the diffusion step.\n    # 3. Momentum-corrected z0t: EMA of clean-image estimates across steps.\n    # 4. Dynamic thresholding (Imagen-style): clamp predicted clean latents.\n\n    def __init__(self, solver_config, model_key=\"runwayml/stable-diffusion-v1-5\",\n                 device=None, **kwargs):\n        super().__init__(solver_config, model_key, device, **kwargs)\n        self.rescale_phi = 0.7\n        self.momentum_beta = 0.12\n        self.thresh_pct = 0.995\n\n    @torch.autocast(device_type='cuda', dtype=torch.float16)\n    def sample(self, cfg_guidance=7.5, prompt=[\"\",\"\"], callback_fn=None, **kwargs):\n        uc, c = self.get_text_embed(null_prompt=prompt[0], prompt=prompt[1])\n        zt = self.initialize_latent().requires_grad_()\n        z0t_prev = None\n        pbar = tqdm(self.scheduler.timesteps, desc=\"CFG++R\")\n        for step, t in enumerate(pbar):\n            at = self.alpha(t)\n            at_prev = self.alpha(t - self.skip)\n            with torch.no_grad():\n                noise_uc, noise_c = self.predict_noise(zt, t, uc, c)\n                noise_cfg = noise_uc + cfg_guidance * (noise_c - noise_uc)\n                # Rescaled CFG: match std of conditional prediction\n                std_cfg = noise_cfg.std(dim=[1,2,3], keepdim=True) + 1e-8\n                std_c = noise_c.std(dim=[1,2,3], keepdim=True)\n                noise_rescaled = noise_cfg * (std_c / std_cfg)\n                noise_pred = self.rescale_phi * noise_rescaled + (1 - self.rescale_phi) * noise_cfg\n            # Tweedie: predict clean image\n            z0t = (zt - (1-at).sqrt() * noise_pred) / at.sqrt()\n            # Dynamic thresholding\n            s = torch.quantile(z0t.abs().reshape(z0t.shape[0], -1), self.thresh_pct, dim=1)\n            s = torch.clamp(s, min=1.0)[:, None, None, None]\n            z0t = z0t.clamp(-s, s) / s\n            # Momentum correction: EMA of z0t (decays with progress)\n            if z0t_prev is not None:\n                progress = step / len(self.scheduler.timesteps)\n                beta = self.momentum_beta * (1.0 - progress)\n                z0t = (1 - beta) * z0t + beta * z0t_prev\n            z0t_prev = z0t.detach().clone()\n            # CFG++ renoising: unconditional noise for stable ODE trajectory\n            zt = at_prev.sqrt() * z0t + (1-at_prev).sqrt() * noise_uc\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': z0t.detach(), 'zt': zt.detach(),\n                                    'decode': self.decode}\n                callback_kwargs = callback_fn(step, t, callback_kwargs)\n                z0t = callback_kwargs[\"z0t\"]\n                zt = callback_kwargs[\"zt\"]\n\n        img = self.decode(z0t)\n        img = (img / 2 + 0.5).clamp(0, 1)\n        return img.detach().cpu()\n"
7}

1{
2  "content": "@register_solver(\"ddim_cfg++\")\nclass BaseDDIMCFGpp(SDXL):\n    # CFG++ with Rescaled Guidance, Momentum Correction & Dynamic Thresholding.\n    # Same algorithmic innovations as the SD1.5 version, adapted for SDXL.\n    rescale_phi = 0.7\n    momentum_beta = 0.12\n    thresh_pct = 0.995\n\n    def reverse_process(self, null_prompt_embeds, prompt_embeds, cfg_guidance,\n                        add_cond_kwargs, shape=(1024, 1024), callback_fn=None, **kwargs):\n        zt = self.initialize_latent(size=(1, 4, shape[1] // self.vae_scale_factor,\n                                          shape[0] // self.vae_scale_factor))\n        z0t_prev = None\n        pbar = tqdm(self.scheduler.timesteps.int(), desc='SDXL-CFG++R')\n        for step, t in enumerate(pbar):\n            next_t = t - self.skip\n            at = self.scheduler.alphas_cumprod[t]\n            at_next = self.scheduler.alphas_cumprod[next_t]\n            with torch.no_grad():\n                noise_uc, noise_c = self.predict_noise(zt, t, null_prompt_embeds,\n                                                        prompt_embeds, add_cond_kwargs)\n                noise_cfg = noise_uc + cfg_guidance * (noise_c - noise_uc)\n                # Rescaled CFG\n                std_cfg = noise_cfg.std(dim=[1,2,3], keepdim=True) + 1e-8\n                std_c = noise_c.std(dim=[1,2,3], keepdim=True)\n                noise_rescaled = noise_cfg * (std_c / std_cfg)\n                noise_pred = self.rescale_phi * noise_rescaled + (1 - self.rescale_phi) * noise_cfg\n            # Tweedie\n            z0t = (zt - (1-at).sqrt() * noise_pred) / at.sqrt()\n            # Dynamic thresholding\n            s = torch.quantile(z0t.abs().reshape(z0t.shape[0], -1), self.thresh_pct, dim=1)\n            s = torch.clamp(s, min=1.0)[:, None, None, None]\n            z0t = z0t.clamp(-s, s) / s\n            # Momentum correction\n            if z0t_prev is not None:\n                progress = step / len(self.scheduler.timesteps)\n                beta = self.momentum_beta * (1.0 - progress)\n                z0t = (1 - beta) * z0t + beta * z0t_prev\n            z0t_prev = z0t.detach().clone()\n            # CFG++ renoising with unconditional noise\n            zt = at_next.sqrt() * z0t + (1-at_next).sqrt() * noise_uc\n            if callback_fn is not None:\n                callback_kwargs = {'z0t': z0t.detach(), 'zt': zt.detach(),\n                                    'decode': self.decode}\n                callback_kwargs = callback_fn(step, t, callback_kwargs)\n                z0t = callback_kwargs[\"z0t\"]\n                zt = callback_kwargs[\"zt\"]\n        return z0t\n\n\n",
3  "end_line": 755,
4  "filename": "CFGpp-main/latent_sdxl.py",
5  "op": "replace",
6  "start_line": 713
7}