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.
This commit is contained in:
Yoland Yan
2025-03-02 12:36:47 -08:00
parent f03ece18f2
commit bfc2f177e8

View File

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