Fix ruff errors

This commit is contained in:
Yoland Yan 2025-03-02 11:44:41 -08:00
parent 225a196dae
commit 2cd3c8a2fb
2 changed files with 16 additions and 20 deletions

View File

@ -1,10 +1,9 @@
import datetime import datetime
import io
import json import json
import math import math
import os import os
import logging
import matplotlib.pyplot as plt
import numpy as np import numpy as np
import safetensors import safetensors
import torch import torch
@ -17,7 +16,6 @@ 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
from nodes import LoadImage
class TrainSampler(comfy.samplers.Sampler): class TrainSampler(comfy.samplers.Sampler):
@ -42,9 +40,9 @@ class TrainSampler(comfy.samplers.Sampler):
loss = self.loss_fn(denoised, latent.clone()) loss = self.loss_fn(denoised, latent.clone())
except RuntimeError as e: except RuntimeError as e:
if "does not require grad and does not have a grad_fn" in str(e): if "does not require grad and does not have a grad_fn" in str(e):
print("WARNING: This is likely due to the model is loaded in inference mode.") logging.info("WARNING: This is likely due to the model is loaded in inference mode.")
loss.backward() loss.backward()
print(f"Current Training Loss: {loss.item():.6f}") logging.info(f"Current Training Loss: {loss.item():.6f}")
if self.loss_callback: if self.loss_callback:
self.loss_callback(loss.item()) self.loss_callback(loss.item())
@ -99,7 +97,7 @@ def load_and_process_images(image_files, input_dir, resize_method="None"):
torch.Tensor: Batch of processed images torch.Tensor: Batch of processed images
""" """
if not image_files: if not image_files:
raise ValueError(f"No valid images found in input") raise ValueError("No valid images found in input")
output_images = [] output_images = []
w, h = None, None w, h = None, None
@ -406,9 +404,7 @@ class TrainLoraNode:
) )
else: else:
if existing_lora != "[None]": if existing_lora != "[None]":
print( logging.info(f"Warning: No existing weights found for {lora_up_key} or {lora_down_key}")
f"Warning: No existing weights found for {lora_up_key} or {lora_down_key}"
)
# Initialize new weights # Initialize new weights
lora_down = torch.nn.Parameter( lora_down = torch.nn.Parameter(
torch.zeros( torch.zeros(

View File

@ -9,11 +9,11 @@ def mock_folder_structure():
# Create a nested folder structure # Create a nested folder structure
folders = [ folders = [
"folder1", "folder1",
os.path.join("folder1", "subfolder1"), "folder1/subfolder1",
os.path.join("folder1", "subfolder2"), "folder1/subfolder2",
"folder2", "folder2",
os.path.join("folder2", "deep"), "folder2/deep",
os.path.join("folder2", "deep", "nested"), "folder2/deep/nested",
"empty_folder" "empty_folder"
] ]