Update src/loss.py

src/loss.py

@@ -1,6 +1,7 @@
 _A=None
 import torch
 from tqdm import tqdm
+import output
 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):
@@ -44,3 +45,14 @@ class SchedulerWrapper:
 	def load_loss_params(A):B,C,D=torch.load(A.loss_params_path,map_location='cpu');A.loss_model=LossSchedulerModel(C,D);A.loss_scheduler=LossScheduler(B,A.loss_model)
 	def load_loss_params_path(A, path):B,C,D=torch.load(path,map_location='cpu');A.loss_model=LossSchedulerModel(C,D);A.loss_scheduler=LossScheduler(B,A.loss_model)
 	def prepare_loss(A,num_accelerate_steps=15):A.load_loss_params()
+def get_instance(device):
+    compress = output.Decoder().to(device).requires_grad_(False)
+    compress.load_state_dict(torch.load("taesdxl_decoder.pth", map_location=device, weights_only=True)) # taesdxl_decoder.pth
+    return compress
+@torch.no_grad()
+def hook_pipe(pipe, compress, mul, sub, scaling_factor):
+    def compress_machine(magic, *args, **kwargs):
+        magic = magic.float().mul(scaling_factor) #pipe.vae.config.scaling_factor
+        out_magic = compress(magic).mul_(mul).sub_(sub).cpu()#.mul_(1.2).sub_(0.75).cpu()
+        return (out_magic, )
+    pipe.vae.decode = compress_machine