Diffusion Plugin

api.py

@@ -47,11 +47,9 @@
     prompts,
     tools,
     batched_prompts,
-    diffusion,
     recipes,
     users,
 )
-
 import torch
 
 try:
@@ -194,7 +192,6 @@ async def validation_exception_handler(request, exc):
 app.include_router(recipes.router)
 app.include_router(batched_prompts.router)
 app.include_router(fastchat_openai_api.router)
-app.include_router(diffusion.router)
 app.include_router(get_xmlrpc_router())
 app.include_router(get_trainer_xmlrpc_router())
 

test/plugins/plugin.schema.json

@@ -15,7 +15,8 @@
         "generator",
         "eval",
         "exporter",
-        "rag"
+        "rag",
+        "diffusion"
       ]
     },
     "version": { "type": "string" },

transformerlab/db/jobs.py

@@ -18,6 +18,7 @@
     "UNDEFINED",
     "GENERATE",
     "INSTALL_RECIPE_DEPS",
+    "DIFFUSION",
 ]
 
 

transformerlab/plugin_sdk/transformerlab/plugin.py

@@ -359,7 +359,7 @@ def add_to_job_data(self, key: str, value):
             cursor = self.db.execute("SELECT job_data FROM job WHERE id = ?", (self.id,))
             row = cursor.fetchone()
             cursor.close()
-            
+
             job_data = {}
             if row and row[0] is not None:
                 data = row[0]
@@ -368,7 +368,7 @@ def add_to_job_data(self, key: str, value):
                     try:
                         # Try to parse as JSON
                         job_data = json.loads(data)
-                        
+
                         # Check if the result is still a string (double-encoded JSON)
                         if isinstance(job_data, str):
                             # Try to parse again
@@ -380,18 +380,18 @@ def add_to_job_data(self, key: str, value):
                     job_data = data
                 else:
                     job_data = {}
-            
+
             # Update the key - handle different value types
             # if isinstance(value, str):
             #     # Try to parse as JSON, if that fails store as string
             #     try:
             #         job_data[key] = json.loads(value)
-                # except (json.JSONDecodeError, TypeError):
-                #     job_data[key] = value
+            # except (json.JSONDecodeError, TypeError):
+            #     job_data[key] = value
             # else:
             # Store value as-is (dict, list, number, bool, etc.)
             job_data[key] = value
-            
+
             # Save back as JSON
             self.db.execute(
                 "UPDATE job SET job_data = ? WHERE id = ?",
@@ -409,7 +409,7 @@ def update_job_data(self, key: str, value):
             cursor = self.db.execute("SELECT job_data FROM job WHERE id = ?", (self.id,))
             row = cursor.fetchone()
             cursor.close()
-            
+
             job_data = {}
             if row and row[0] is not None:
                 data = row[0]
@@ -418,7 +418,7 @@ def update_job_data(self, key: str, value):
                     try:
                         # Try to parse as JSON
                         job_data = json.loads(data)
-                        
+
                         # Check if the result is still a string (double-encoded JSON)
                         if isinstance(job_data, str):
                             # Try to parse again
@@ -430,7 +430,7 @@ def update_job_data(self, key: str, value):
                     job_data = data
                 else:
                     job_data = {}
-            
+
             # Update the key - handle different value types
             if isinstance(value, str):
                 # Try to parse as JSON, if that fails store as string
@@ -441,7 +441,7 @@ def update_job_data(self, key: str, value):
             else:
                 # Store value as-is (dict, list, number, bool, etc.)
                 job_data[key] = value
-            
+
             # Save back as JSON
             self.db.execute(
                 "UPDATE job SET job_data = ? WHERE id = ?",

transformerlab/plugin_sdk/transformerlab/sdk/v1/diffusion.py

@@ -0,0 +1,36 @@
+from transformerlab.sdk.v1.tlab_plugin import TLabPlugin
+
+
+class DiffusionTLabPlugin(TLabPlugin):
+    """Enhanced decorator class for TransformerLab diffusion plugins"""
+
+    def __init__(self):
+        super().__init__()
+        self._parser.add_argument("--run_name", default="diffused", type=str, help="Name for the diffusion output")
+        self._parser.add_argument("--experiment_name", default="default", type=str, help="Name of the experiment")
+        self._parser.add_argument("--diffusion_model", default="local", type=str, help="Diffusion model to use")
+        self._parser.add_argument("--model", type=str, default="")
+
+        self.tlab_plugin_type = "diffusion"
+
+    def _ensure_args_parsed(self):
+        if not self._args_parsed:
+            args, unknown_args = self._parser.parse_known_args()
+            for key, value in vars(args).items():
+                self.params[key] = value
+            self._parse_unknown_args(unknown_args)
+            self._args_parsed = True
+
+    def _parse_unknown_args(self, unknown_args):
+        key = None
+        for arg in unknown_args:
+            if arg.startswith("--"):
+                key = arg.lstrip("-")
+                self.params[key] = True
+            elif key:
+                self.params[key] = arg
+                key = None
+
+
+# Global instance (like tlab_gen in generate.py)
+tlab_diffusion = DiffusionTLabPlugin()

transformerlab/plugins/image_diffusion/index.json

@@ -0,0 +1,16 @@
+{
+  "name": "Image Diffusion",
+  "uniqueId": "image_diffusion",
+  "description": "Generate images in the Diffusion tab using this plugin",
+  "plugin-format": "python",
+  "type": "diffusion",
+  "version": "0.0.1",
+  "git": "",
+  "url": "",
+  "files": ["main.py", "diffusion_worker.py", "setup.sh"],
+  "supported_hardware_architectures": ["cuda", "amd"],
+  "_dataset": false,
+  "setup-script": "setup.sh",
+  "parameters": {},
+  "parameters_ui": {}
+}