From bfc2f177e812bb95dae4a5064af771d90ca0d31a Mon Sep 17 00:00:00 2001
From: Yoland Yan <4950057+yoland68@users.noreply.github.com>
Date: Sun, 2 Mar 2025 12:36:47 -0800
Subject: [PATCH] Refactor import statements in nodes_train.py

Reorganized and cleaned up import statements, removing unused imports and adding specific module imports for better clarity and organization.
---
 comfy_extras/nodes_train.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py
index 9b777f217..8a0c5f435 100644
--- a/comfy_extras/nodes_train.py
+++ b/comfy_extras/nodes_train.py
@@ -1,8 +1,8 @@
 import datetime
 import json
+import logging
 import math
 import os
-import logging
 
 import numpy as np
 import safetensors
@@ -10,13 +10,15 @@ import torch
 from PIL import Image, ImageDraw, ImageFont
 from PIL.PngImagePlugin import PngInfo
 
-import comfy
-import comfy_extras
+import comfy.samplers
+import comfy.utils
+import comfy_extras.nodes_custom_sampler
 import folder_paths
 import node_helpers
 from comfy.cli_args import args
 from comfy.comfy_types.node_typing import IO
 
+
 class TrainSampler(comfy.samplers.Sampler):
 
     def __init__(self, loss_fn, optimizer, loss_callback=None):