Update src/loss.py

src/loss.py

@@ -74,30 +74,53 @@ class LossScheduler:
 		A.model=model
 		A.init_noise_sigma=1.
 		A.order=1
-		A.optimizer = torch.optim.SGD([A.model.scale_factor, A.model.wx, A.model.we], lr=lr)
-
+		# A.optimizer = torch.optim.SGD([A.model.scale_factor, A.model.wx, A.model.we], lr=lr)
+		A.optimizer = torch.optim.SGD([A.model.wx, A.model.we], lr=lr)
+		A.model.train()
 	@staticmethod
 	def load(path):A,B,C,E=torch.load(path,map_location='cpu');D=LossSchedulerModel(B,C,scale_factor=E);return LossScheduler(A,D)
 	def save(A,path):B,C,D,E=A.timesteps,A.model.wx,A.model.we,A.model.scale_factor;torch.save((B,C,D,E),path)
 	def set_timesteps(A,num_inference_steps,device='cuda'):B=device;A.xT=_A;A.e_prev=[];A.t_prev=-1;A.model=A.model.to(B);A.timesteps=A.timesteps.to(B)
 	def scale_model_input(A,sample,*B,**C):return sample
-	def step(self,model_output,timestep,sample,*D,**E):
+	def step_orig(self,model_output,timestep,sample,*D,**E):
 		A=self;B=A.timesteps.tolist().index(timestep);assert A.t_prev==-1 or B==A.t_prev+1
 		if A.t_prev==-1:A.xT=sample
 		A.e_prev.append(model_output);C=A.model(B,A.xT,A.e_prev)
 		if B+1==len(A.timesteps):A.xT=_A;A.e_prev=[];A.t_prev=-1
 		else:A.t_prev=B
 		return C,
+	def step(self,model_output,timestep,sample,*D,**E):
+		global image_count
+		global count
+		if timestep==B.timesteps[0]: print("resetting"); image_count+=1;count=0;
+		A=self;B=A.timesteps.tolist().index(timestep);assert A.t_prev==-1 or B==A.t_prev+1; count+=1
+		if A.t_prev==-1:A.xT=sample
+		A.e_prev.append(model_output)
+		if timestep==B.timesteps[-1]:
+			with torch.enable_grad():
+				orig_output = torch.load(f"/home/mbhat/edge-maxxing/miner/miner/latents/latent_oig_{image_count}_20.pth")
+				A.optimizer.zero_grad()
+				C=A.model(B,A.xT,A.e_prev)
+				L=torch.nn.functional.mse_loss(C, orig_output)
+				L.backward()
+				A.optimizer.step()
+		if B+1==len(A.timesteps):
+			A.xT=_A;A.e_prev=[];
+			A.t_prev=-1
+			# A.save("/home/mbhat/weights.pth")    
+		else:A.t_prev=B
+		return C,
 	def step_other(self,model_output,timestep,sample,*D,**E):
 		A=self;B=A.timesteps.tolist().index(timestep);assert A.t_prev==-1 or B==A.t_prev+1; count+=1
 		if A.t_prev==-1:A.xT=sample
 		A.e_prev.append(model_output)
-		with torch.enable_grad():
-			A.optimizer.zero_grad()
-			C=A.model(B,A.xT,A.e_prev)
-			L=torch.nn.functional.mse_loss(C, model_output)
-			L.backward()
-			A.optimizer.step()
+		if timestep==A.timesteps.tolist()[-1]:
+			with torch.enable_grad():
+				A.optimizer.zero_grad()
+				C=A.model(B,A.xT,A.e_prev)
+				L=torch.nn.functional.mse_loss(C, model_output)
+				L.backward()
+				A.optimizer.step()
 		if B+1==len(A.timesteps):
 			A.xT=_A;A.e_prev=[];
 			A.t_prev=-1
@@ -115,11 +138,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");
-		if timestep: print(timestep); 
+		
 		if B.loss_scheduler is _A:
 			C=B.scheduler.step(E,A,D,**F);A=A.tolist();
 			if A not in B.catch_x:B.catch_x[A]=[];B.catch_e[A]=[];B.catch_x_[A]=[]