Update src/loss.py

src/loss.py

@@ -1,6 +1,9 @@
 _A=None
+count=0
+image_count=0
 import torch
 from tqdm import tqdm
+import numpy as np
 class LossSchedulerModel(torch.nn.Module):
 	def __init__(A,wx,we):super(LossSchedulerModel,A).__init__();assert len(wx.shape)==1 and len(we.shape)==2;B=wx.shape[0];assert B==we.shape[0]and B==we.shape[1];A.register_parameter('wx',torch.nn.Parameter(wx));A.register_parameter('we',torch.nn.Parameter(we))
 	def forward(A,t,xT,e_prev):
@@ -29,6 +32,8 @@ class SchedulerWrapper:
 		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
 		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
 	def step(B,model_output,timestep,sample,**F):
+		global image_count
+		global count
 		D=sample;E=model_output;A=timestep
 		if timestep==B.timesteps[0]: print("resetting"); image_count+=1;count=0; torch.save(model_output, f"/home/mbhat/edge-maxxing/miner/miner/latents/latent_rob_{image_count}_0.pth")
 		if True: