diff --git a/PixArt/config.py b/PixArt/config.py
index e11e26a..43082a1 100644
--- a/PixArt/config.py
+++ b/PixArt/config.py
@@ -1,12 +1,10 @@
-"""
-Model config and setting logic
-"""
 import math
 import logging
 
 import comfy.supported_models_base
 import comfy.supported_models
 import comfy.latent_formats
+import comfy.model_base
 
 from .model.pixart import PixArt
 from .model.pixartms import PixArtMS
diff --git a/PixArt/loader.py b/PixArt/loader.py
index d63b1d8..9692a71 100644
--- a/PixArt/loader.py
+++ b/PixArt/loader.py
@@ -1,28 +1,14 @@
-import comfy.supported_models_base
-import comfy.latent_formats
-import comfy.model_detection
-import comfy.model_patcher
-import comfy.model_base
-import comfy.utils
-import comfy.conds
 import logging
-import torch
 
+import comfy.utils
+import comfy.model_base
+import comfy.model_patcher
+import comfy.model_detection
 from comfy import model_management
+
 from .diffusers_convert import convert_state_dict
 from .config import model_config_from_unet
 
-class PixArtModel(comfy.model_base.BaseModel):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-    
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-
-        for name in ["width", "height", "aspect_ratio", "img_hw"]: # TODO: remove last one
-            out[name] = comfy.conds.CONDRegular(torch.tensor(name))
-        return out
-
 def load_pixart_state_dict(sd, model_options={}):
     # prefix / format
     sd = sd.get("model", sd) # ref ckpt
diff --git a/nodes/pixart.py b/nodes/pixart.py
deleted file mode 100644
index e69de29..0000000