diff --git a/modules/editor/shadertools/converter/functions/customfunction.h b/modules/editor/shadertools/converter/functions/customfunction.h
index 56cbef9fc..382acd5ca 100644
--- a/modules/editor/shadertools/converter/functions/customfunction.h
+++ b/modules/editor/shadertools/converter/functions/customfunction.h
@@ -114,7 +114,7 @@ class CustomFunction : public ShaderNode {
     int32_t build(QString &code, QStack<QString> &stack, const AbstractNodeGraph::Link &link, int32_t &depth, int32_t &type) override {
         if(m_position == -1) {
             if(!m_func.isEmpty()) {
-                static_cast<ShaderGraph *>(m_graph)->addFunction(m_funcName, m_func);
+                static_cast<ShaderGraph *>(m_graph)->addFragmentFunction(m_funcName, m_func);
 
                 if(link.oport->m_type != MetaType::INVALID) {
                     type = link.oport->m_type;
diff --git a/modules/editor/shadertools/converter/functions/matrixoperations.h b/modules/editor/shadertools/converter/functions/matrixoperations.h
index efc0860b3..3221475d6 100644
--- a/modules/editor/shadertools/converter/functions/matrixoperations.h
+++ b/modules/editor/shadertools/converter/functions/matrixoperations.h
@@ -170,7 +170,7 @@ class ExtractPosition : public ShaderNode {
                           "    return vec4(m[3][0], m[3][1], m[3][2], m[3][3]);\n"
                           "}\n";
 
-            static_cast<ShaderGraph *>(m_graph)->addFunction("ExtractPosition", func);
+            static_cast<ShaderGraph *>(m_graph)->addFragmentFunction("ExtractPosition", func);
 
             QStringList arguments = getArguments(code, stack, depth, type);
 
diff --git a/modules/editor/shadertools/converter/shadergraph.cpp b/modules/editor/shadertools/converter/shadergraph.cpp
index b9de63391..370c07e57 100644
--- a/modules/editor/shadertools/converter/shadergraph.cpp
+++ b/modules/editor/shadertools/converter/shadergraph.cpp
@@ -213,6 +213,8 @@ ShaderGraph::ShaderGraph() :
     Split::registerClassFactory(Engine::resourceSystem());
     Swizzle::registerClassFactory(Engine::resourceSystem());
 
+    CustomFunction::registerClassFactory(Engine::resourceSystem());
+
     // Common
     NodeGroup::registerClassFactory(Engine::resourceSystem());
 
@@ -307,7 +309,7 @@ GraphNode *ShaderGraph::nodeCreate(const TString &type, int &index) {
         return node;
     } else { // Self exposed function
         if(!type.isEmpty()) {
-            CustomFunction *function = new CustomFunction();
+            CustomFunction *function = Engine::objectCreate<CustomFunction>(type);
             function->exposeFunction(m_exposedFunctions[type]);
             function->setGraph(this);
 
@@ -404,12 +406,19 @@ bool ShaderGraph::buildGraph(GraphNode *node) {
     setPragma("uniforms", layout);
 
     // Functions
-    TString functions;
-    for(const auto &it : m_functions) {
-        functions += it.second + '\n';
+    TString vertexFunctions;
+    for(const auto &it : m_vertexFunctions) {
+        vertexFunctions += it.second + '\n';
+    }
+
+    setPragma("vertexFunctions", vertexFunctions);
+
+    TString fragmentFunctions;
+    for(const auto &it : m_fragmentFunctions) {
+        fragmentFunctions += it.second + '\n';
     }
 
-    setPragma("functions", functions);
+    setPragma("fragmentFunctions", fragmentFunctions);
 
     return true;
 }
@@ -558,10 +567,17 @@ void ShaderGraph::addUniform(const TString &name, uint8_t type, const Variant &v
     m_uniforms.push_back({name, type, 1, value});
 }
 
-void ShaderGraph::addFunction(const TString &name, TString &code) {
-    auto it = m_functions.find(name);
-    if(it == m_functions.end()) {
-        m_functions[name] = code;
+void ShaderGraph::addVertexFunction(const TString &name, TString &code) {
+    auto it = m_vertexFunctions.find(name);
+    if(it == m_vertexFunctions.end()) {
+        m_vertexFunctions[name] = code;
+    }
+}
+
+void ShaderGraph::addFragmentFunction(const TString &name, TString &code) {
+    auto it = m_fragmentFunctions.find(name);
+    if(it == m_fragmentFunctions.end()) {
+        m_fragmentFunctions[name] = code;
     }
 }
 
@@ -673,7 +689,8 @@ TString ShaderGraph::buildFrom(GraphNode *node, Stage stage) {
 void ShaderGraph::cleanup() {
     m_textures.clear();
     m_uniforms.clear();
-    m_functions.clear();
+    m_vertexFunctions.clear();
+    m_fragmentFunctions.clear();
     m_pragmas.clear();
 }
 
diff --git a/modules/editor/shadertools/converter/shadergraph.h b/modules/editor/shadertools/converter/shadergraph.h
index 81da724c2..43d8e450d 100644
--- a/modules/editor/shadertools/converter/shadergraph.h
+++ b/modules/editor/shadertools/converter/shadergraph.h
@@ -35,7 +35,8 @@ class ShaderGraph : public AbstractNodeGraph {
 
     void addUniform(const TString &name, uint8_t type, const Variant &value);
 
-    void addFunction(const TString &name, TString &code);
+    void addVertexFunction(const TString &name, TString &code);
+    void addFragmentFunction(const TString &name, TString &code);
 
     StringList nodeList() const override;
 
@@ -116,7 +117,8 @@ class ShaderGraph : public AbstractNodeGraph {
 
     std::list<std::pair<TString, int32_t>> m_textures;
 
-    std::map<TString, TString> m_functions;
+    std::map<TString, TString> m_vertexFunctions;
+    std::map<TString, TString> m_fragmentFunctions;
 
     std::map<TString, TString> m_exposedFunctions;
 
diff --git a/modules/editor/shadertools/shaders/Billboard.vert b/modules/editor/shadertools/shaders/Billboard.vert
index 36f22e925..da0ac1887 100644
--- a/modules/editor/shadertools/shaders/Billboard.vert
+++ b/modules/editor/shadertools/shaders/Billboard.vert
@@ -23,7 +23,7 @@ layout(location = 7) flat out vec4 _objectId;
 layout(location = 8) flat out int _instanceOffset;
 layout(location = 9) out mat4 _modelView;
 
-#pragma functions
+#pragma vertexFunctions
 
 void main(void) {
 #pragma offset
diff --git a/modules/editor/shadertools/shaders/Shader.frag b/modules/editor/shadertools/shaders/Shader.frag
index 40323530a..49b2200ff 100644
--- a/modules/editor/shadertools/shaders/Shader.frag
+++ b/modules/editor/shadertools/shaders/Shader.frag
@@ -32,7 +32,7 @@ layout(location = 0) out vec4 gbuffer0;
 
 #include "Functions.h"
 
-#pragma functions
+#pragma fragmentFunctions
 
 void main(void) {
 #pragma instance
diff --git a/modules/editor/shadertools/shaders/Skinned.vert b/modules/editor/shadertools/shaders/Skinned.vert
index e07dbcce9..988f9dd71 100644
--- a/modules/editor/shadertools/shaders/Skinned.vert
+++ b/modules/editor/shadertools/shaders/Skinned.vert
@@ -31,7 +31,7 @@ layout(location = 7) flat out vec4 _objectId;
 layout(location = 8) flat out int _instanceOffset;
 layout(location = 9) out mat4 _modelView;
 
-#pragma functions
+#pragma vertexFunctions
 
 void main(void) {
     #pragma skinOffset
diff --git a/modules/editor/shadertools/shaders/Static.vert b/modules/editor/shadertools/shaders/Static.vert
index c5fe405f8..bbbe6a5e0 100644
--- a/modules/editor/shadertools/shaders/Static.vert
+++ b/modules/editor/shadertools/shaders/Static.vert
@@ -28,7 +28,7 @@ layout(location = 7) flat out vec4 _objectId;
 layout(location = 8) flat out int _instanceOffset;
 layout(location = 9) out mat4 _modelView;
 
-#pragma functions
+#pragma vertexFunctions
 
 void main(void) {
 #pragma offset