Skip to content

Commit 2d55a5f

Browse files
mashb1tmashb1t
authored
feat: add support for playground v2.5 (#3073)
* feat: add support for playground v2.5 * feat: add preset for playground v2.5 * feat: change URL to mashb1t * feat: optimize playground v2.5 preset
1 parent cb24c68 commit 2d55a5f

File tree

11 files changed

+233
-30
lines changed

11 files changed

+233
-30
lines changed

‎args_manager.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,4 @@
11
import ldm_patched.modules.args_parser as args_parser
2-
import os
3-
4-
from tempfile import gettempdir
52

63
args_parser.parser.add_argument("--share", action='store_true', help="Set whether to share on Gradio.")
74

‎ldm_patched/contrib/external_model_advanced.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ class ModelSamplingContinuousEDM:
108108
@classmethod
109109
def INPUT_TYPES(s):
110110
return {"required": { "model": ("MODEL",),
111-
"sampling": (["v_prediction", "eps"],),
111+
"sampling": (["v_prediction", "edm_playground_v2.5", "eps"],),
112112
"sigma_max": ("FLOAT", {"default": 120.0, "min": 0.0, "max": 1000.0, "step":0.001, "round": False}),
113113
"sigma_min": ("FLOAT", {"default": 0.002, "min": 0.0, "max": 1000.0, "step":0.001, "round": False}),
114114
}}
@@ -121,17 +121,25 @@ def INPUT_TYPES(s):
121121
def patch(self, model, sampling, sigma_max, sigma_min):
122122
m = model.clone()
123123

124+
latent_format = None
125+
sigma_data = 1.0
124126
if sampling == "eps":
125127
sampling_type = ldm_patched.modules.model_sampling.EPS
126128
elif sampling == "v_prediction":
127129
sampling_type = ldm_patched.modules.model_sampling.V_PREDICTION
130+
elif sampling == "edm_playground_v2.5":
131+
sampling_type = ldm_patched.modules.model_sampling.EDM
132+
sigma_data = 0.5
133+
latent_format = ldm_patched.modules.latent_formats.SDXL_Playground_2_5()
128134

129135
class ModelSamplingAdvanced(ldm_patched.modules.model_sampling.ModelSamplingContinuousEDM, sampling_type):
130136
pass
131137

132138
model_sampling = ModelSamplingAdvanced(model.model.model_config)
133-
model_sampling.set_sigma_range(sigma_min, sigma_max)
139+
model_sampling.set_parameters(sigma_min, sigma_max, sigma_data)
134140
m.add_object_patch("model_sampling", model_sampling)
141+
if latent_format is not None:
142+
m.add_object_patch("latent_format", latent_format)
135143
return (m, )
136144

137145
class RescaleCFG:

‎ldm_patched/modules/latent_formats.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import torch
12

23
class LatentFormat:
34
scale_factor = 1.0
@@ -34,6 +35,70 @@ def __init__(self):
3435
]
3536
self.taesd_decoder_name = "taesdxl_decoder"
3637

38+
class SDXL_Playground_2_5(LatentFormat):
39+
def __init__(self):
40+
self.scale_factor = 0.5
41+
self.latents_mean = torch.tensor([-1.6574, 1.886, -1.383, 2.5155]).view(1, 4, 1, 1)
42+
self.latents_std = torch.tensor([8.4927, 5.9022, 6.5498, 5.2299]).view(1, 4, 1, 1)
43+
44+
self.latent_rgb_factors = [
45+
# R G B
46+
[ 0.3920, 0.4054, 0.4549],
47+
[-0.2634, -0.0196, 0.0653],
48+
[ 0.0568, 0.1687, -0.0755],
49+
[-0.3112, -0.2359, -0.2076]
50+
]
51+
self.taesd_decoder_name = "taesdxl_decoder"
52+
53+
def process_in(self, latent):
54+
latents_mean = self.latents_mean.to(latent.device, latent.dtype)
55+
latents_std = self.latents_std.to(latent.device, latent.dtype)
56+
return (latent - latents_mean) * self.scale_factor / latents_std
57+
58+
def process_out(self, latent):
59+
latents_mean = self.latents_mean.to(latent.device, latent.dtype)
60+
latents_std = self.latents_std.to(latent.device, latent.dtype)
61+
return latent * latents_std / self.scale_factor + latents_mean
62+
63+
3764
class SD_X4(LatentFormat):
3865
def __init__(self):
3966
self.scale_factor = 0.08333
67+
self.latent_rgb_factors = [
68+
[-0.2340, -0.3863, -0.3257],
69+
[ 0.0994, 0.0885, -0.0908],
70+
[-0.2833, -0.2349, -0.3741],
71+
[ 0.2523, -0.0055, -0.1651]
72+
]
73+
74+
class SC_Prior(LatentFormat):
75+
def __init__(self):
76+
self.scale_factor = 1.0
77+
self.latent_rgb_factors = [
78+
[-0.0326, -0.0204, -0.0127],
79+
[-0.1592, -0.0427, 0.0216],
80+
[ 0.0873, 0.0638, -0.0020],
81+
[-0.0602, 0.0442, 0.1304],
82+
[ 0.0800, -0.0313, -0.1796],
83+
[-0.0810, -0.0638, -0.1581],
84+
[ 0.1791, 0.1180, 0.0967],
85+
[ 0.0740, 0.1416, 0.0432],
86+
[-0.1745, -0.1888, -0.1373],
87+
[ 0.2412, 0.1577, 0.0928],
88+
[ 0.1908, 0.0998, 0.0682],
89+
[ 0.0209, 0.0365, -0.0092],
90+
[ 0.0448, -0.0650, -0.1728],
91+
[-0.1658, -0.1045, -0.1308],
92+
[ 0.0542, 0.1545, 0.1325],
93+
[-0.0352, -0.1672, -0.2541]
94+
]
95+
96+
class SC_B(LatentFormat):
97+
def __init__(self):
98+
self.scale_factor = 1.0 / 0.43
99+
self.latent_rgb_factors = [
100+
[ 0.1121, 0.2006, 0.1023],
101+
[-0.2093, -0.0222, -0.0195],
102+
[-0.3087, -0.1535, 0.0366],
103+
[ 0.0290, -0.1574, -0.4078]
104+
]

‎ldm_patched/modules/model_sampling.py‎

Lines changed: 80 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import torch
2-
import numpy as np
32
from ldm_patched.ldm.modules.diffusionmodules.util import make_beta_schedule
43
import math
54

@@ -12,12 +11,28 @@ def calculate_denoised(self, sigma, model_output, model_input):
1211
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
1312
return model_input - model_output * sigma
1413

14+
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
15+
if max_denoise:
16+
noise = noise * torch.sqrt(1.0 + sigma ** 2.0)
17+
else:
18+
noise = noise * sigma
19+
20+
noise += latent_image
21+
return noise
22+
23+
def inverse_noise_scaling(self, sigma, latent):
24+
return latent
1525

1626
class V_PREDICTION(EPS):
1727
def calculate_denoised(self, sigma, model_output, model_input):
1828
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
1929
return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) - model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
2030

31+
class EDM(V_PREDICTION):
32+
def calculate_denoised(self, sigma, model_output, model_input):
33+
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
34+
return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) + model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
35+
2136

2237
class ModelSamplingDiscrete(torch.nn.Module):
2338
def __init__(self, model_config=None):
@@ -42,24 +57,23 @@ def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps
4257
else:
4358
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
4459
alphas = 1. - betas
45-
alphas_cumprod = torch.tensor(np.cumprod(alphas, axis=0), dtype=torch.float32)
46-
# alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
60+
alphas_cumprod = torch.cumprod(alphas, dim=0)
4761

4862
timesteps, = betas.shape
4963
self.num_timesteps = int(timesteps)
5064
self.linear_start = linear_start
5165
self.linear_end = linear_end
5266

67+
# self.register_buffer('betas', torch.tensor(betas, dtype=torch.float32))
68+
# self.register_buffer('alphas_cumprod', torch.tensor(alphas_cumprod, dtype=torch.float32))
69+
# self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32))
70+
5371
sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
5472
self.set_sigmas(sigmas)
55-
self.set_alphas_cumprod(alphas_cumprod.float())
5673

5774
def set_sigmas(self, sigmas):
58-
self.register_buffer('sigmas', sigmas)
59-
self.register_buffer('log_sigmas', sigmas.log())
60-
61-
def set_alphas_cumprod(self, alphas_cumprod):
62-
self.register_buffer("alphas_cumprod", alphas_cumprod.float())
75+
self.register_buffer('sigmas', sigmas.float())
76+
self.register_buffer('log_sigmas', sigmas.log().float())
6377

6478
@property
6579
def sigma_min(self):
@@ -94,18 +108,18 @@ def percent_to_sigma(self, percent):
94108
class ModelSamplingContinuousEDM(torch.nn.Module):
95109
def __init__(self, model_config=None):
96110
super().__init__()
97-
self.sigma_data = 1.0
98-
99111
if model_config is not None:
100112
sampling_settings = model_config.sampling_settings
101113
else:
102114
sampling_settings = {}
103115

104116
sigma_min = sampling_settings.get("sigma_min", 0.002)
105117
sigma_max = sampling_settings.get("sigma_max", 120.0)
106-
self.set_sigma_range(sigma_min, sigma_max)
118+
sigma_data = sampling_settings.get("sigma_data", 1.0)
119+
self.set_parameters(sigma_min, sigma_max, sigma_data)
107120

108-
def set_sigma_range(self, sigma_min, sigma_max):
121+
def set_parameters(self, sigma_min, sigma_max, sigma_data):
122+
self.sigma_data = sigma_data
109123
sigmas = torch.linspace(math.log(sigma_min), math.log(sigma_max), 1000).exp()
110124

111125
self.register_buffer('sigmas', sigmas) #for compatibility with some schedulers
@@ -134,3 +148,56 @@ def percent_to_sigma(self, percent):
134148

135149
log_sigma_min = math.log(self.sigma_min)
136150
return math.exp((math.log(self.sigma_max) - log_sigma_min) * percent + log_sigma_min)
151+
152+
class StableCascadeSampling(ModelSamplingDiscrete):
153+
def __init__(self, model_config=None):
154+
super().__init__()
155+
156+
if model_config is not None:
157+
sampling_settings = model_config.sampling_settings
158+
else:
159+
sampling_settings = {}
160+
161+
self.set_parameters(sampling_settings.get("shift", 1.0))
162+
163+
def set_parameters(self, shift=1.0, cosine_s=8e-3):
164+
self.shift = shift
165+
self.cosine_s = torch.tensor(cosine_s)
166+
self._init_alpha_cumprod = torch.cos(self.cosine_s / (1 + self.cosine_s) * torch.pi * 0.5) ** 2
167+
168+
#This part is just for compatibility with some schedulers in the codebase
169+
self.num_timesteps = 10000
170+
sigmas = torch.empty((self.num_timesteps), dtype=torch.float32)
171+
for x in range(self.num_timesteps):
172+
t = (x + 1) / self.num_timesteps
173+
sigmas[x] = self.sigma(t)
174+
175+
self.set_sigmas(sigmas)
176+
177+
def sigma(self, timestep):
178+
alpha_cumprod = (torch.cos((timestep + self.cosine_s) / (1 + self.cosine_s) * torch.pi * 0.5) ** 2 / self._init_alpha_cumprod)
179+
180+
if self.shift != 1.0:
181+
var = alpha_cumprod
182+
logSNR = (var/(1-var)).log()
183+
logSNR += 2 * torch.log(1.0 / torch.tensor(self.shift))
184+
alpha_cumprod = logSNR.sigmoid()
185+
186+
alpha_cumprod = alpha_cumprod.clamp(0.0001, 0.9999)
187+
return ((1 - alpha_cumprod) / alpha_cumprod) ** 0.5
188+
189+
def timestep(self, sigma):
190+
var = 1 / ((sigma * sigma) + 1)
191+
var = var.clamp(0, 1.0)
192+
s, min_var = self.cosine_s.to(var.device), self._init_alpha_cumprod.to(var.device)
193+
t = (((var * min_var) ** 0.5).acos() / (torch.pi * 0.5)) * (1 + s) - s
194+
return t
195+
196+
def percent_to_sigma(self, percent):
197+
if percent <= 0.0:
198+
return 999999999.9
199+
if percent >= 1.0:
200+
return 0.0
201+
202+
percent = 1.0 - percent
203+
return self.sigma(torch.tensor(percent))

‎ldm_patched/modules/samplers.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -523,7 +523,7 @@ def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=N
523523

524524
KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
525525
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
526-
"dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "tcd"]
526+
"dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "tcd", "edm_playground_v2.5"]
527527

528528
class KSAMPLER(Sampler):
529529
def __init__(self, sampler_function, extra_options={}, inpaint_options={}):

‎modules/async_worker.py‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -828,16 +828,33 @@ def handler(async_task):
828828

829829
if scheduler_name in ['lcm', 'tcd']:
830830
final_scheduler_name = 'sgm_uniform'
831-
if pipeline.final_unet is not None:
832-
pipeline.final_unet = core.opModelSamplingDiscrete.patch(
831+
832+
def patch_discrete(unet):
833+
return core.opModelSamplingDiscrete.patch(
833834
pipeline.final_unet,
834835
sampling=scheduler_name,
835836
zsnr=False)[0]
837+
838+
if pipeline.final_unet is not None:
839+
pipeline.final_unet = patch_discrete(pipeline.final_unet)
836840
if pipeline.final_refiner_unet is not None:
837-
pipeline.final_refiner_unet = core.opModelSamplingDiscrete.patch(
838-
pipeline.final_refiner_unet,
841+
pipeline.final_refiner_unet = patch_discrete(pipeline.final_refiner_unet)
842+
print(f'Using {scheduler_name} scheduler.')
843+
elif scheduler_name == 'edm_playground_v2.5':
844+
final_scheduler_name = 'karras'
845+
846+
def patch_edm(unet):
847+
return core.opModelSamplingContinuousEDM.patch(
848+
unet,
839849
sampling=scheduler_name,
840-
zsnr=False)[0]
850+
sigma_max=120.0,
851+
sigma_min=0.002)[0]
852+
853+
if pipeline.final_unet is not None:
854+
pipeline.final_unet = patch_edm(pipeline.final_unet)
855+
if pipeline.final_refiner_unet is not None:
856+
pipeline.final_refiner_unet = patch_edm(pipeline.final_refiner_unet)
857+
841858
print(f'Using {scheduler_name} scheduler.')
842859

843860
async_task.yields.append(['preview', (flags.preparation_step_count, 'Moving model to GPU ...', None)])

‎modules/core.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,7 @@
2121
from modules.util import get_file_from_folder_list
2222
from ldm_patched.modules.lora import model_lora_keys_unet, model_lora_keys_clip
2323
from modules.config import path_embeddings
24-
from ldm_patched.contrib.external_model_advanced import ModelSamplingDiscrete
25-
24+
from ldm_patched.contrib.external_model_advanced import ModelSamplingDiscrete, ModelSamplingContinuousEDM
2625

2726
opEmptyLatentImage = EmptyLatentImage()
2827
opVAEDecode = VAEDecode()
@@ -32,6 +31,7 @@
3231
opControlNetApplyAdvanced = ControlNetApplyAdvanced()
3332
opFreeU = FreeU_V2()
3433
opModelSamplingDiscrete = ModelSamplingDiscrete()
34+
opModelSamplingContinuousEDM = ModelSamplingContinuousEDM()
3535

3636

3737
class StableDiffusionModel:

‎modules/flags.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,7 @@
4848

4949
KSAMPLER_NAMES = list(KSAMPLER.keys())
5050

51-
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo",
52-
"align_your_steps", "tcd"]
51+
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo", "align_your_steps", "tcd", "edm_playground_v2.5"]
5352
SAMPLER_NAMES = KSAMPLER_NAMES + list(SAMPLER_EXTRA.keys())
5453

5554
sampler_list = SAMPLER_NAMES

‎modules/patch_precision.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,6 @@ def patched_register_schedule(self, given_betas=None, beta_schedule="linear", ti
5151
self.linear_end = linear_end
5252
sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32)
5353
self.set_sigmas(sigmas)
54-
alphas_cumprod = torch.tensor(alphas_cumprod, dtype=torch.float32)
55-
self.set_alphas_cumprod(alphas_cumprod)
5654
return
5755

5856

‎presets/.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,6 @@
22
!anime.json
33
!default.json
44
!lcm.json
5+
!playground_v2.5.json
56
!realistic.json
67
!sai.json

0 commit comments

Comments
 (0)