Update worker_runpod.py

worker_runpod.py

@@ -6,27 +6,26 @@ import numpy as np
 
 import nodes
 from nodes import NODE_CLASS_MAPPINGS
+from comfy_extras import nodes_flux
 
-with torch.inference_mode():
-    DualCLIPLoader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
-    UNETLoader = NODE_CLASS_MAPPINGS["UNETLoader"]()
-    VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
-    ControlNetLoader = NODE_CLASS_MAPPINGS["ControlNetLoader"]()
+DualCLIPLoader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
+UNETLoader = NODE_CLASS_MAPPINGS["UNETLoader"]()
+VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
+ControlNetLoader = NODE_CLASS_MAPPINGS["ControlNetLoader"]()
 
-    with torch.inference_mode():
-        clip = DualCLIPLoader.load_clip("t5xxl_fp16.safetensors", "clip_l.safetensors", "flux")[0]
-        unet = UNETLoader.load_unet("flux1-dev-fp8.safetensors", "fp8_e4m3fn")[0]
-        vae = VAELoader.load_vae("ae.safetensors")[0]
-        controlnet = ControlNetLoader.load_controlnet("controlnet.safetensors")[0]
+FluxGuidance = nodes_flux.NODE_CLASS_MAPPINGS["FluxGuidance"]()
+VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
+VAEDecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
+VAEEncode = NODE_CLASS_MAPPINGS["VAEEncode"]()
+ControlNetApplyAdvanced = NODE_CLASS_MAPPINGS["ControlNetApplyAdvanced"]()
+ImageScaleBy = NODE_CLASS_MAPPINGS["ImageScaleBy"]()
+LoadImage = NODE_CLASS_MAPPINGS["LoadImage"]()
 
-    from comfy_extras import nodes_flux
-    FluxGuidance = nodes_flux.NODE_CLASS_MAPPINGS["FluxGuidance"]()
-    VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
-    VAEDecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
-    VAEEncode = NODE_CLASS_MAPPINGS["VAEEncode"]()
-    ControlNetApplyAdvanced = NODE_CLASS_MAPPINGS["ControlNetApplyAdvanced"]()
-    ImageScaleBy = NODE_CLASS_MAPPINGS["ImageScaleBy"]()
-    LoadImage = NODE_CLASS_MAPPINGS["LoadImage"]()
+with torch.inference_mode():
+    clip = DualCLIPLoader.load_clip("t5xxl_fp16.safetensors", "clip_l.safetensors", "flux")[0]
+    unet = UNETLoader.load_unet("flux1-dev-fp8.safetensors", "fp8_e4m3fn")[0]
+    vae = VAELoader.load_vae("ae.safetensors")[0]
+    controlnet = ControlNetLoader.load_controlnet("controlnet.safetensors")[0]
 
 def closestNumber(n, m):
     q = int(n / m)