nooooooo se

app.py

@@ -46,8 +46,8 @@ config = Config(
     tokenizer=CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
 )
 
-# Load models with SE blocks enabled
-config.models = model_loader.load_models(str(model_file), device, use_se=True)
+# Load models
+config.models = model_loader.load_models(str(model_file), device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024

src/model_loader.py

@@ -6,7 +6,7 @@ from .diffusion import Diffusion
 from . import model_converter
 import torch
 
-def load_models(ckpt_path, device, use_se=False):
+def load_models(ckpt_path, device):
     state_dict = model_converter.load_from_standard_weights(ckpt_path, device)
 
     encoder = VAE_Encoder().to(device)
@@ -15,19 +15,9 @@ def load_models(ckpt_path, device, use_se=False):
     decoder = VAE_Decoder().to(device)
     decoder.load_state_dict(state_dict['decoder'], strict=True)
 
-    # Initialize diffusion model with SE blocks disabled for loading pre-trained weights
-    diffusion = Diffusion(use_se=False).to(device)
+    # Initialize diffusion model
+    diffusion = Diffusion().to(device)
     diffusion.load_state_dict(state_dict['diffusion'], strict=True)
-    
-    # If SE blocks are requested, reinitialize the model with them
-    if use_se:
-        diffusion = Diffusion(use_se=True).to(device)
-        # Copy the weights from the loaded model
-        with torch.no_grad():
-            for name, param in diffusion.named_parameters():
-                if 'se' not in name:  # Skip SE block parameters
-                    if name in state_dict['diffusion']:
-                        param.copy_(state_dict['diffusion'][name])
 
     clip = CLIP().to(device)
     clip.load_state_dict(state_dict['clip'], strict=True)