LiruiZhao commited on
Commit
db29364
1 Parent(s): b7b1d93
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -88,12 +88,16 @@ class CompVisDenoiser(K.external.CompVisDenoiser):
88
  print("c_out.device:", c_out.device)
89
  print("sigma.device:", sigma.device)
90
  # eps_0, eps_1 = self.get_eps(input_0 * c_in, input_1 * c_in, self.sigma_to_t(sigma), **kwargs)
91
- eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma.cpu()).cuda(), **kwargs)
92
 
93
  return input_0 + eps_0 * c_out, eps_1
94
 
95
  def to_d(x, sigma, denoised):
96
  """Converts a denoiser output to a Karras ODE derivative."""
 
 
 
 
97
  return (x - denoised) / append_dims(sigma, x.ndim)
98
 
99
  def default_noise_sampler(x):
@@ -206,7 +210,7 @@ def generate(
206
  print("uncond['c_crossattn'][0].device:", uncond["c_crossattn"][0].device)
207
  print("uncond['c_concat'][0].device:", uncond["c_concat"][0].device)
208
 
209
- sigmas = model_wrap.get_sigmas(steps)
210
 
211
  extra_args = {
212
  "cond": cond,
 
88
  print("c_out.device:", c_out.device)
89
  print("sigma.device:", sigma.device)
90
  # eps_0, eps_1 = self.get_eps(input_0 * c_in, input_1 * c_in, self.sigma_to_t(sigma), **kwargs)
91
+ eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma.cpu().float()).cuda(), **kwargs)
92
 
93
  return input_0 + eps_0 * c_out, eps_1
94
 
95
  def to_d(x, sigma, denoised):
96
  """Converts a denoiser output to a Karras ODE derivative."""
97
+
98
+ print(x.device)
99
+ print(denoised.device)
100
+ print(sigma.device)
101
  return (x - denoised) / append_dims(sigma, x.ndim)
102
 
103
  def default_noise_sampler(x):
 
210
  print("uncond['c_crossattn'][0].device:", uncond["c_crossattn"][0].device)
211
  print("uncond['c_concat'][0].device:", uncond["c_concat"][0].device)
212
 
213
+ sigmas = model_wrap.get_sigmas(steps).to(model.device)
214
 
215
  extra_args = {
216
  "cond": cond,