initial commit of ComfyUI node

__init__.py

@@ -0,0 +1,10 @@
+import os
+import sys
+sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
+
+from node import Eden_LoRa_trainer
+
+NODE_CLASS_MAPPINGS = {
+    "Eden_LoRa_trainer": Eden_LoRa_trainer,
+}
+

node.py

@@ -0,0 +1,113 @@
+
+import os
+import shutil
+import tarfile
+import json
+import time
+import random
+import torch
+import numpy as np
+import pandas as pd
+
+from dotenv import load_dotenv
+from main import train
+
+from trainer.preprocess import preprocess
+from trainer.models import pretrained_models
+from trainer.config import TrainingConfig
+from trainer.utils.io import clean_filename
+from trainer.utils.utils import seed_everything
+
+class Eden_LoRa_trainer:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                    "training_images_folder_path": ("STRING", {"default": "."}),
+                    "sd_model_version": (["sdxl", "sd15"], ),
+                    "seed": ("INT", {"default": 0, "min": 0, "max": 100000}),
+                    "resolution": ("INT", {"default": 512, "min": 256, "max": 768}),
+                    "train_batch_size": ("INT", {"default": 4, "min": 1, "max": 8}),
+                    "max_train_steps":  ("INT", {"default": 400, "min": 50, "max": 1000}),
+                    "ti_lr":   ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 0.01, "step": 0.0001}),
+                    "unet_lr": ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 0.01, "step": 0.0001}),
+                    "lora_rank": ("INT", {"default": 16, "min": 1, "max": 64}),
+                    "use_dora": ("BOOLEAN", {"default": False}),
+                    "n_tokens": ("INT", {"default": 2, "min": 1, "max": 3}),
+                }
+        }
+
+    CATEGORY = "Eden 🌱"
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "train_lora"
+
+    def train_lora(self, training_images_folder_path,
+            concept_mode = "style",
+            sd_model_version = "sdxl",
+            seed = 0,
+            resolution = 521,
+            train_batch_size = 4,
+            max_train_steps = 400,
+            ti_lr = 0.001,
+            unet_lr = 0.001,
+            lora_rank = 16,
+            use_dora = False,
+            n_tokens = 2
+            ):
+        
+        print("Starting new training job...")
+
+        config = TrainingConfig(
+            name="test",
+            lora_training_urls=folder,
+            concept_mode=concept_mode,
+            sd_model_version=sd_model_version,
+            seed=seed,
+            resolution=resolution,
+            train_batch_size=train_batch_size,
+            max_train_steps=max_train_steps,
+            checkpointing_steps=10000,
+            ti_lr=ti_lr,
+            unet_lr=unet_lr,
+            lora_rank=lora_rank,
+            use_dora=use_dora,
+            caption_model="blip",
+            n_tokens=n_tokens,
+            verbose=True,
+            debug=True,
+        )
+        
+        train_generator = train(config=config)
+
+        while True:
+            try:
+                progress_f = next(train_generator)
+            except StopIteration as e:
+                config, output_save_dir = e.value  # Capture the return value
+                break
+
+        validation_grid_img_path = os.path.join(output_save_dir, "validation_grid.jpg")
+        out_path = f"{clean_filename(name)}_eden_concept_lora_{int(time.time())}.tar"
+        directory = cogPath(output_save_dir)
+
+        with tarfile.open(out_path, "w") as tar:
+            print("Adding files to tar...")
+            for file_path in directory.rglob("*"):
+                print(file_path)
+                arcname = file_path.relative_to(directory)
+                tar.add(file_path, arcname=arcname)
+            
+            # Add instructions README:
+            tar.add("instructions_README.md", arcname="README.md")
+            tar.add("comfyUI_workflow_lora_txt2img.json", arcname="comfyUI_workflow_lora_txt2img.json")
+            if sd_model_version == "sd15":
+                tar.add("comfyUI_workflow_lora_adiff.json", arcname="comfyUI_workflow_lora_adiff.json")
+
+        attributes = {}
+        attributes['grid_prompts'] = config.training_attributes["validation_prompts"]
+        attributes['job_time_seconds'] = config.job_time
+
+        print(f"LORA training finished in {config.job_time:.1f} seconds")
+        print(f"Returning {out_path}")
+
+        return (out_path,)

requirements.txt

@@ -1,17 +1,17 @@
-torch==2.1.0
-torchvision==0.16.0
-transformers==4.38.1
-diffusers==0.27.2
-ujson==5.9.0
-scipy==1.12.0
-peft==0.10.0
-invisible-watermark==0.2.0
-accelerate==0.27.2
+torch>=2.1.0
+torchvision>=0.16.0
+transformers>=4.38.1
+diffusers>=0.27.2
+ujson>=5.9.0
+scipy>=1.12.0
+peft>=0.10.0
+invisible-watermark>=0.2.0
 pandas==2.2.1
-numpy==1.26.4
+numpy>=1.26.4
 opencv-python>=4.1.0.25
-mediapipe==0.10.11
-openai==1.14.0
+mediapipe>=0.10.11
+openai>=1.14.0
 python-dotenv
 prodigyopt
-omegaconf
+omegaconf
+ujson