Update src/loss.py
Browse files- src/loss.py +1 -1
src/loss.py
CHANGED
|
@@ -25,7 +25,7 @@ class LossScheduler:
|
|
| 25 |
class SchedulerWrapper:
|
| 26 |
def __init__(A,scheduler,loss_params_path='loss_model_new_final_ulta_small.pth'):A.scheduler=scheduler;A.catch_x,A.catch_e,A.catch_x_={},{},{};A.loss_scheduler=_A;A.loss_params_path=loss_params_path
|
| 27 |
def set_timesteps(A,num_inference_steps,**C):
|
| 28 |
-
D=
|
| 29 |
if A.loss_scheduler is _A:B=A.scheduler.set_timesteps(D,**C);A.timesteps=A.scheduler.timesteps;A.init_noise_sigma=A.scheduler.init_noise_sigma;A.order=A.scheduler.order;return B
|
| 30 |
else:B=A.loss_scheduler.set_timesteps(D,**C);A.timesteps=A.loss_scheduler.timesteps;A.init_noise_sigma=A.scheduler.init_noise_sigma;A.order=A.scheduler.order;return B
|
| 31 |
def step(B,model_output,timestep,sample,**F):
|
|
|
|
| 25 |
class SchedulerWrapper:
|
| 26 |
def __init__(A,scheduler,loss_params_path='loss_model_new_final_ulta_small.pth'):A.scheduler=scheduler;A.catch_x,A.catch_e,A.catch_x_={},{},{};A.loss_scheduler=_A;A.loss_params_path=loss_params_path
|
| 27 |
def set_timesteps(A,num_inference_steps,**C):
|
| 28 |
+
D=10
|
| 29 |
if A.loss_scheduler is _A:B=A.scheduler.set_timesteps(D,**C);A.timesteps=A.scheduler.timesteps;A.init_noise_sigma=A.scheduler.init_noise_sigma;A.order=A.scheduler.order;return B
|
| 30 |
else:B=A.loss_scheduler.set_timesteps(D,**C);A.timesteps=A.loss_scheduler.timesteps;A.init_noise_sigma=A.scheduler.init_noise_sigma;A.order=A.scheduler.order;return B
|
| 31 |
def step(B,model_output,timestep,sample,**F):
|