From 56b22f3858a56b09d5de34412b9d4e1f7d8a8237 Mon Sep 17 00:00:00 2001
From: Andrew Suter-Morris <andrew@prenav.com>
Date: Wed, 19 Feb 2025 14:54:55 -0700
Subject: [PATCH] clean up

---
 README.md        |   1 +
 __init__.py      |   6 +++
 try_off_nodes.py | 110 ++++++++++++-----------------------------------
 3 files changed, 35 insertions(+), 82 deletions(-)

diff --git a/README.md b/README.md
index 49b0d02..e9159aa 100644
--- a/README.md
+++ b/README.md
@@ -20,3 +20,4 @@ To use, use the 8BitQuantized.json. You can remove the Quanitization step and it
 - Optimize, optimize, optimize.
 - TryOn
 - Custom VAE, Text Encoders, etc.
+- Use LoRA
diff --git a/__init__.py b/__init__.py
index abcfe19..c29e83b 100644
--- a/__init__.py
+++ b/__init__.py
@@ -4,6 +4,8 @@ from .try_off_nodes import (  # noqa
     TryOffModelNode,
     TryOffQuantizerNode,
     TryOffRunNode,
+    TryOnOffModelNode,
+    FluxFillPipelineNode2
 )
 
 # Node Registrations
@@ -13,6 +15,8 @@ NODE_CLASS_MAPPINGS = {
     "TryOffRunNode": TryOffRunNode,
     "TryOffQuantizerNode": TryOffQuantizerNode,
     "TryOffFluxFillPipelineNode": FluxFillPipelineNode,
+    "TryOnOffModelNode": TryOnOffModelNode,
+    "TryOffFluxFillPipelineNode2": FluxFillPipelineNode2
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -21,4 +25,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "TryOffRunNode": "Run TryOff Inference [ComfyUI-Flux-TryOff]",
     "TryOffQuantizerNode": "TryOff Quantizer [ComfyUI-Flux-TryOff]",
     "TryOffFluxFillPipelineNode": "FluxFill Pipeline Loader [ComfyUI-Flux-TryOff]",
+    "TryOnOffModelNode": "TryOn or TryOff Model Loader [ComfyUI-Flux-TryOff]",
+    "TryOffFluxFillPipelineNode2": "FluxFill Pipeline Loader [ComfyUI-Flux-TryOff]",
 }
diff --git a/try_off_nodes.py b/try_off_nodes.py
index 52d0e11..5ed748e 100644
--- a/try_off_nodes.py
+++ b/try_off_nodes.py
@@ -279,18 +279,20 @@ class TryOnOffModelNode:
             ).to(device)
         return (model,)
     
+
 class FluxFillPipelineNode2:
     @classmethod
     def INPUT_TYPES(cls):  # noqa: N802
         return {
             "required": {
                 "transformer": ("MODEL",),
+                "vae": ("MODEL",), 
+                "text_encoder": ("MODEL",),
+                "tokenizer": ("TOKENIZER",),
+                "text_encoder_2": ("MODEL",),
+                "tokenizer_2": ("TOKENIZER",),
                 "device": (device_list,),
-            },
-            "optional": {
-                "transformers_config": ("transformers_config",),
-                "diffusers_config": ("diffusers_config",),
-            },
+            }
         }
 
     CATEGORY = "Models"
@@ -298,86 +300,29 @@ class FluxFillPipelineNode2:
     FUNCTION = "load_pipeline"
 
     def load_pipeline(
-        self, transformer, device, transformers_config=None, diffusers_config=None
+        self,
+        transformer,
+        vae,
+        text_encoder,
+        tokenizer,
+        text_encoder_2,
+        tokenizer_2,
+        device,
     ):
-        if transformers_config:
-            tokenizer = CLIPTokenizer.from_pretrained(
-                "openai/clip-vit-large-patch14",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-                quantization_config=transformers_config,
-            )
-            tokenizer_2 = T5TokenizerFast.from_pretrained(
-                "XLabs-AI/xflux_text_encoders",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-                quantization_config=transformers_config,
-            )
-            text_encoder = CLIPTextModel.from_pretrained(
-                "openai/clip-vit-large-patch14",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-                quantization_config=transformers_config,
-            )
-            text_encoder_2 = T5EncoderModel.from_pretrained(
-                "XLabs-AI/xflux_text_encoders",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-                quantization_config=transformers_config,
-            )
-        else:
-            tokenizer = CLIPTokenizer.from_pretrained(
-                "openai/clip-vit-large-patch14",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-            )
-            tokenizer_2 = T5TokenizerFast.from_pretrained(
-                "XLabs-AI/xflux_text_encoders",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-            )
-            text_encoder = CLIPTextModel.from_pretrained(
-                "openai/clip-vit-large-patch14",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-            )
-            text_encoder_2 = T5EncoderModel.from_pretrained(
-                "XLabs-AI/xflux_text_encoders",
-                cache_dir=encoders_dir,
-                torch_dtype=dtype,
-            )
-
+        # Initialize scheduler independently (assuming this is still needed)
         scheduler = FlowMatchEulerDiscreteScheduler()
 
-        if diffusers_config:
-            vae = AutoencoderTiny.from_pretrained(
-                "madebyollin/taef1",
-                cache_dir=vae_dir,
-                torch_dtype=dtype,
-                quantization_config=diffusers_config,
-            )
-            pipeline = FluxFillPipeline(
-                scheduler=scheduler,
-                vae=vae,
-                text_encoder=text_encoder,
-                tokenizer=tokenizer,
-                text_encoder_2=text_encoder_2,
-                tokenizer_2=tokenizer_2,
-                transformer=transformer
-            )
-        else:
-            vae = AutoencoderTiny.from_pretrained(
-                "madebyollin/taef1", cache_dir=vae_dir, torch_dtype=dtype
-            )
-            pipeline = FluxFillPipeline(
-                scheduler=scheduler,
-                vae=vae,
-                text_encoder=text_encoder,
-                tokenizer=tokenizer,
-                text_encoder_2=text_encoder_2,
-                tokenizer_2=tokenizer_2,
-                transformer=transformer,
-            )
+        # Construct the FluxFillPipeline using the supplied components
+        pipeline = FluxFillPipeline(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+        )
+
         pipeline.to(device)
         pipeline.enable_model_cpu_offload()
         pipeline.transformer.to(dtype)
@@ -385,6 +330,7 @@ class FluxFillPipelineNode2:
         return (pipeline,)
 
 
+
 # TryOffRun Node
 class TryOffRunNode:
     @classmethod