Update src/pipeline.py

src/pipeline.py

@@ -13,7 +13,6 @@ import numpy as np
 import os
 import pickle
 from tqdm import tqdm
-import math
 
 import torchvision.models as models
 import torch.nn.functional as F
@@ -69,39 +68,22 @@ class LossSchedulerModel(torch.nn.Module):
 
         return result #.to(xT.dtype)
 
-
-
-        
 class LossSchedulerModel(torch.nn.Module):
-    def __init__(A,wx=None,we=None,wb=None, timesteps=13, input_shape=[13, 13]):
+    def __init__(A,wx,we,wb=None):
         super(LossSchedulerModel,A).__init__()
-        if wx is None:
-            A.register_parameter('wx',torch.nn.Parameter(A.init1d(timesteps)))
-        else:
-            A.register_parameter('wx',torch.nn.Parameter(wx))
-        if wb is None:
-            A.register_parameter('we',torch.nn.Parameter(A.init2d(A.wx.shape[0], A.wx.shape[0])))
-        else:
-            A.register_parameter('we',torch.nn.Parameter(we))
-        if wb is None:
-            A.register_parameter('wb',torch.nn.Parameter(A.init1d(A.wx.shape[0])))
-        else:
-            A.register_parameter('wb',torch.nn.Parameter(wb))
         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))
+        if wb is None:
+            A.register_parameter('wb',torch.nn.Parameter(torch.randn(wx.shape)))
+        else: 
+            A.register_parameter('wb',torch.nn.Parameter(wb)
         
     def forward(A,t,xT,e_prev):
         B=e_prev;assert t-len(B)+1==0;C=xT*A.wx[t]
         for(D,E)in zip(B,A.we[t]):C+=D*E
         C += A.wb[t]
         return C.to(xT.dtype)
-
-    def init1d(A, point):
-        stdv = 1. / math.sqrt(point)
-        return torch.empty(point).uniform_(-stdv, stdv)
-
-    def init2d(A, rows, cols):
-        stdv = 1. / math.sqrt(rows)
-        return torch.empty(rows, cols).uniform_(-stdv, stdv)
         
 class LossScheduler:
 	def __init__(A,timesteps,model):A.timesteps=timesteps;A.model=model;A.init_noise_sigma=1.;A.order=1
@@ -363,7 +345,7 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
     alpha = 0.60
     # loss_fn = PerceptualLoss()
     # Training loop
-    num_epochs = 2500
+    num_epochs = 2000
     
     for epoch in range(num_epochs):
         total_loss = 0
@@ -425,19 +407,19 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
     timesteps_sel = torch.tensor(timesteps_sel)
     # scheduler = LossScheduler(timesteps=timesteps_sel, model=loss_model)
     # scheduler.save(f'../optimized_loss_params_{name}.pth')
-    torch.save(pipeline.scheduler.loss_model.state_dict(), f'../large_loss_params_{name}.pth')
+    torch.save(pipeline.scheduler.loss_model.state_dict(), f'../large_12_loss_params_{name}.pth')
     print("model saved")
 
     # Update the pipeline with the new parameters
     # pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(orig_scheduler.config))
     # pipeline.scheduler.load_loss_params_path(f'../optimized_loss_params_{name}.pth')  # Reload the updated parameters
-    model = LossSchedulerModel(input_shape=(1, 4, 128, 128))
-    model.load_state_dict(torch.load(f'../large_loss_params_{name}.pth', map_location=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')))
+    model = LossSchedulerModel(pipeline.scheduler.loss_model.wx, pipeline.scheduler.loss_model.we) #input_shape=(1, 4, 128, 128)
+    model.load_state_dict(torch.load(f'../large_12_loss_params_{name}.pth', map_location=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')))
     model.eval()  
     pipeline.scheduler.loss_scheduler.model = model
     pipeline.scheduler.loss_model = model
     output_images = pipeline(prompts[0], output_type="pil", num_inference_steps=len(selected_indices), generator=generator, guidance_scale=5.0).images
-    output_images[0].save(f"../large_final_image_{name}.png")
+    output_images[0].save(f"../large_12_final_image_{name}.png")
 
     # # import pdb; pdb.set_trace()
     # simialrity = torch.nn.CosineSimilarity(dim=1)(output_images[0].flatten(1), target_images[0].flatten(1)).mean()
@@ -466,4 +448,4 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
         guidance_rescale = 0.0,
         callback_on_step_end=callback_dynamic_cfg,
         callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],
-    ).images[0]
+    ).images[0]