wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-23 17:02:38 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Update samplers.py

Browse Source
This commit is contained in:
hlky 2023-05-15 19:35:51 +01:00
parent abba81d8a6
commit 6c1c6232eb
1 changed files with 1 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
comfy/samplers.py
View File

@ -1,7 +1,6 @@
from .k_diffusion import sampling as k_diffusion_sampling from .k_diffusion import sampling as k_diffusion_sampling
from .k_diffusion import external as k_diffusion_external from .k_diffusion import external as k_diffusion_external
from .extra_samplers import uni_pc from .extra_samplers import uni_pc
import os
import torch import torch
import contextlib import contextlib
from diffusers import LMSDiscreteScheduler from diffusers import LMSDiscreteScheduler
@ -9,7 +8,7 @@ from comfy import model_management
from .ldm.models.diffusion.ddim import DDIMSampler from .ldm.models.diffusion.ddim import DDIMSampler
from .ldm.modules.diffusionmodules.util import make_ddim_timesteps from .ldm.modules.diffusionmodules.util import make_ddim_timesteps
import math import math
from aitemplate.compiler import Model
def lcm(a, b): #TODO: eventually replace by math.lcm (added in python3.9) def lcm(a, b): #TODO: eventually replace by math.lcm (added in python3.9)
return abs(a*b) // math.gcd(a, b) return abs(a*b) // math.gcd(a, b)
@ -532,25 +531,12 @@ class AITemplateModelWrapper:
shape[1] = height shape[1] = height
shape[2] = width shape[2] = width
ys.append(torch.empty(shape).cuda().half()) ys.append(torch.empty(shape).cuda().half())
# print(inputs["input0"].shape)
# print(inputs["input1"].shape)
# print(inputs["input2"].shape)
# print(ys)
self.unet_ait_exe.run_with_tensors(inputs, ys, graph_mode=False) self.unet_ait_exe.run_with_tensors(inputs, ys, graph_mode=False)
noise_pred = ys[0].permute((0, 3, 1, 2)).float() noise_pred = ys[0].permute((0, 3, 1, 2)).float()
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
return noise_pred return noise_pred
def init_ait_module(
model_name,
workdir,
):
mod = Model(os.path.join(workdir, model_name, "test.so"))
return mod
class KSampler: class KSampler:
SCHEDULERS = ["normal", "karras", "simple", "ddim_uniform"] SCHEDULERS = ["normal", "karras", "simple", "ddim_uniform"]
SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral", SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",