fix no init

hydra_node/config.py

@@ -12,6 +12,22 @@ from sentry_sdk import capture_exception
 from sentry_sdk.integrations.threading import ThreadingIntegration
 from hydra_node.models import StableDiffusionModel
 
+def no_init(loading_code):
+    def dummy(self):
+        return
+    
+    modules = [torch.nn.Linear, torch.nn.Embedding, torch.nn.LayerNorm]
+    original = {}
+    for mod in modules:
+        original[mod] = mod.reset_parameters
+        mod.reset_parameters = dummy
+    
+    result = loading_code()
+    for mod in modules:
+        mod.reset_parameters = original[mod]
+    
+    return result
+
 def init_config_model():
     config = DotMap()
     config.dtype = os.getenv("DTYPE", "float16")
@@ -81,7 +97,7 @@ def init_config_model():
     load_time = time.time()
 
     try:
-        model = StableDiffusionModel(config)
+        model = no_init(lambda: StableDiffusionModel(config))
     except Exception as e:
         logger.error(f"Failed to load model: {str(e)}")
         capture_exception(e)
@@ -97,7 +113,7 @@ def init_config_model():
 
     f = open("/tmp/health_readiness", "w")
     f.close()
-    
+
     time_load = time.time() - load_time
     logger.info(f"Models loaded in {time_load:.2f}s")
 

hydra_node/models.py

@@ -14,27 +14,11 @@ from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.models.diffusion.plms import PLMSSampler
 import time
 
-def no_init(loading_code):
-    def dummy(self):
-        return
-    
-    modules = [torch.nn.Linear, torch.nn.Embedding, torch.nn.LayerNorm]
-    original = {}
-    for mod in modules:
-        original[mod] = mod.reset_parameters
-        mod.reset_parameters = dummy
-    
-    result = loading_code()
-    for mod in modules:
-        mod.reset_parameters = original[mod]
-    
-    return result
-
 class StableDiffusionModel(nn.Module):
     def __init__(self, config):
         nn.Module.__init__(self)
         self.config = config
-        model, model_config = no_init(lambda: self.from_folder(config.model_path))
+        model, model_config = self.from_folder(config.model_path)
         if config.dtype == "float16":
             typex = torch.float16
         else: