Diffusion Plugin #421

Sourenm · 2025-07-08T15:59:28Z

No description provided.

transformerlab/routers/experiment/diffusion.py

codecov · 2025-07-08T19:36:09Z

Codecov Report

Attention: Patch coverage is 22.66667% with 928 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
transformerlab/plugins/image_diffusion/main.py	1.58%	620 Missing ⚠️
transformerlab/routers/experiment/diffusion.py	50.87%	226 Missing and 26 partials ⚠️
transformerlab/shared/shared.py	0.00%	56 Missing ⚠️

📢 Thoughts on this report? Let us know!

transformerlab/routers/experiment/diffusion.py

deep1401 · 2025-07-11T22:04:19Z

api.py

    recipes,
    users,
 )
-
+from transformerlab.routers.experiment import diffusion


We do not import it like this, look at the routers/experiment/experiment.py file inside and see how other routes under experiment are imported and follow that structure

deep1401 · 2025-07-11T22:06:31Z

transformerlab/plugin_sdk/transformerlab/plugin.py

@@ -452,6 +452,10 @@ def set_job_completion_status(
            self.add_to_job_data("completion_status", completion_status)
            self.add_to_job_data("completion_details", completion_details)

+            # Update the job status field if there's a failure
+            if completion_status == "failed":


Is this a duplicate change?
Maybe remove this as we do the same on the next line

the duplicate is removed.

deep1401 · 2025-07-11T22:08:36Z

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

+                self.params[key] = arg
+                key = None
+
+    def save_generated_images(self, metadata: dict = None, suffix: str = None):


Not sure what this is being used for? If its a dataset then it is valid but if you're trying to save a generated image with this then this might be wrong?

I removed this function and the get_output_file_path function but kept diffusion.py since it might come in handy for future diffusion plugins

deep1401 · 2025-07-11T22:09:39Z

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

+    def get_output_file_path(self, suffix="", dir_only=False):
+        self._ensure_args_parsed()
+
+        workspace_dir = os.environ.get("_TFL_WORKSPACE_DIR", "./")


Instead of this could you import WORKSPACE_DIR from transformerlab.plugin and use that?

solved (comment above)

deep1401 · 2025-07-11T22:11:22Z

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

+        print(f"[DIFFUSION] Metadata saved to: {metadata_file}")
+        return metadata_file
+
+    def get_output_file_path(self, suffix="", dir_only=False):


I understand why we add this here but if this is the exact same with only a difference of "diffusion" and "datasets", then it would be wiser to move this out to the base TLabPlugin class and add a parameter so GenTLab and DiffusionTLab can both use

solved (comment above)

deep1401 · 2025-07-11T22:14:32Z

transformerlab/plugins/image_diffusion/main.py

@@ -720,7 +483,11 @@ def get_pipeline(
        # Load LoRA adaptor if provided - same code for local and HF Hub!
        if adaptor and adaptor.strip():
            try:
-                adaptor_dir = os.path.join(os.environ.get("_TFL_WORKSPACE_DIR"), "adaptors", secure_filename(model))
+                adaptor_dir = os.path.join(
+                    os.environ.get("_TFL_WORKSPACE_DIR"),


This should be imported from transformerlab.plugin as an earlier comment

deep1401 · 2025-07-11T22:16:35Z

transformerlab/plugins/image_diffusion/main.py

@@ -1097,7 +1160,22 @@ def run_pipe():
            generation_start = time.time()
            log_print("Starting image generation...")

+            progress_queue = queue.Queue()
+
+            async def monitor_progress():


What is this used for? I thought we do progress update inside the callback

deep1401 · 2025-07-11T22:17:07Z

transformerlab/plugins/image_diffusion/main.py


+def parse_args():


We should not be parsing any arg inside a plugin

deep1401 · 2025-07-11T22:17:32Z

transformerlab/plugins/image_diffusion/main.py

+    parser.add_argument("--plugin_dir", type=str, required=True)
+    parser.add_argument("--job_id", type=str, required=True)
+    parser.add_argument("--experiment_name", type=str, required=True)
+    parser.add_argument("--run_name", type=str, default="diffused")


I guess this is stale code and you need to remove it?

deep1401 · 2025-07-11T22:19:31Z

transformerlab/shared/shared.py

+        config = job_config.get("config", {})
+
+        # Convert base64 images to files and update config
+        base64_fields = {


Why do you we need to add this to config? Why cant the plugin directly fetch them?

We need to save these images because if the input/mask image is passed as a huge base64 string via CLI args, it hits a wall where Linux has a system-imposed limit on command-line length (typically around 128 KB to 256 KB total across all args and env vars).

deep1401 · 2025-07-11T22:21:37Z

transformerlab/plugins/image_diffusion/main.py

-    return {"controlnets": models}
+def main():
+    print("Starting image diffusion...", flush=True)
+    asyncio.run(async_main())


Why do we need to do this? And if you do need this, try to see how we have a async_job_wrapper

deep1401 · 2025-07-14T22:28:38Z

transformerlab/plugins/image_diffusion/main.py

+@tlab_diffusion.async_job_wrapper(progress_start=0, progress_end=100)
+async def diffusion_generate_job():
+    job_config = {
+        "plugin": tlab_diffusion.params.get("plugin", "image_diffusion"),


We probably don't need this?

deep1401 · 2025-07-14T22:29:03Z

transformerlab/plugins/image_diffusion/main.py

+
+@tlab_diffusion.async_job_wrapper(progress_start=0, progress_end=100)
+async def diffusion_generate_job():
+    job_config = {


Why aren't we directly using the tlab_diffusion.params instead of maintaining this?

deep1401 · 2025-07-14T22:29:49Z

transformerlab/plugins/image_diffusion/main.py

+        "mask_image_path": "mask_image",
+    }
+
+    for path_key, base64_key in image_path_keys.items():


Just to clarify here on previous comment, this makes the images back to base64 for processing by diffusers right?

Yes that's correct

deep1401 · 2025-07-14T22:32:14Z

transformerlab/plugins/image_diffusion/setup.sh

Just a note: we left this blank for now because all files are being installed by the original requirements

deep1401 · 2025-07-14T22:32:43Z

transformerlab/plugins/image_diffusion/main.py

+import asyncio
+import threading
+import gc
+from transformerlab.plugin import WORKSPACE_DIR


Maybe we standardize using this throughtout the file?

deep1401 · 2025-07-14T22:33:39Z

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

+        self._parser.add_argument("--diffusion_model", default="local", type=str, help="Diffusion model to use")
+        self._parser.add_argument("--model", type=str, default="")
+        self._parser.add_argument(
+            "--diffusion_type", default="txt2img", type=str, help="Type of diffusion task (txt2img, img2img, etc.)"


Do we use this param somewhere?

deep1401 · 2025-07-14T22:34:47Z

transformerlab/shared/shared.py

Adding it here because its probably not included in the diff, now that we include diffusion_worker in the plugin, we should not have a copy in the shared folder

deep1401 · 2025-07-14T22:38:38Z

transformerlab/shared/shared.py

+        except Exception as e:
+            await db_jobs.job_update_status(job_id, "FAILED")
+            print(f"[DIFFUSION] Job {job_id} execution error: {e}")
+            return {"status": "error", "job_id": job_id, "message": "Diffusion job failed"}


For error returning, there are some cases like when using a SDXL model and the height and width are not multiples of 8, this will error out but we want to show a correct error on why it errored out. The current implementation prints out the error but maybe we could send out the HTTPResponse for the error somehow and show that via notifications?

Additionally, we will also need to stream output from this job, as it is shown for other jobs, maybe we include a button where we show the progress to see the raw output of the generation which uses the same ViewOutputStreamingModal on the app side?

Either we do the HTTPResponse thing or we make use of how the errors are reported on /generate right now in the except block and we show the errors from the job within there as well if a job has failed status

deep1401 · 2025-07-14T22:45:45Z

transformerlab/routers/experiment/diffusion.py

+
+router = APIRouter(prefix="/diffusion", tags=["diffusion"])
+
+UNIVERSIAL_GENERATION_ID = str(uuid.uuid4())


I don't know how we reached to using this but this should not be used here.
This is why we have the /generate_id route which generates a generation_id beforehand and you either use that or one provided with the /generate but the UNIVERSIAL_GENERATION_ID should not be used

…ab-api into add/change-diffusion-to-plugin

transformerlab/routers/experiment/diffusion.py

deep1401

This is almost ready to merge, just a couple of small changes to ensure we still support the plugins we created surrounding this still work

deep1401 · 2025-07-16T14:23:13Z

transformerlab/routers/experiment/diffusion.py

+            raise HTTPException(status_code=400, detail="num_images must be between 1 and 10")
+
+        # Validate diffusion type
+        if request.plugin == "image_diffusion":


A minor change here, once you check the name of request.plugin, you should also check whether it is installed or not, refer to the plugins.py router on how this is checked. But if the plugin is not installed, you return an error. This will fix the support issue we have with the dataset_imagegen plugin then

Maybe after this in dataset_imagegen, we show the error from the request incase generation fails because of the plugin not being installed

deep1401 · 2025-07-16T14:24:21Z

transformerlab/routers/experiment/diffusion.py

+    diffusion_logger.info(message)
+
+
+class DiffusionOutputCapture:


We do not use this anymore so maybe we remove it?

dadmobile · 2025-07-22T02:00:38Z

This worked for me!

convert diffusion's /generate into a job

8e8f5ef

github-advanced-security bot found potential problems Jul 8, 2025

View reviewed changes

transformerlab/routers/experiment/diffusion.py Fixed Show fixed Hide fixed

transformerlab/routers/experiment/diffusion.py Fixed Show fixed Hide fixed

Sourena and others added 8 commits July 8, 2025 09:14

security fix

871492e

fix for test_diffusion.py

c82b46d

fix for test_diffusion.py

af77eb9

fix for test_diffusion.py

f6939a2

fix for test_diffusion.py

d8e7aad

Merge branch 'main' into add/change-diffusion-to-plugin

5b92b46

yet another fix for test_diffusion.py

82101f1

fix status_code assertion

5214903

Sourenm and others added 9 commits July 9, 2025 07:49

Merge branch 'main' into add/change-diffusion-to-plugin

32a4fed

convert image diffusion to a plugin

9f8b7bb

Merge branch 'main' into add/change-diffusion-to-plugin

c7e0e24

fix for test_diffusion.py

60b26dc

add "diffusion" to allowed types in plugin.schema.json

c54ab30

Merge branch 'main' into add/change-diffusion-to-plugin

0ed3a44

progress bar for image_diffusion plugin

60fa0a2

Merge branch 'main' into add/change-diffusion-to-plugin

beb2c33

support polling for diffusion job

3b8ea74

github-advanced-security bot found potential problems Jul 10, 2025

View reviewed changes

transformerlab/routers/experiment/diffusion.py Fixed Show fixed Hide fixed

transformerlab/routers/experiment/diffusion.py Fixed Show fixed Hide fixed

Sourena and others added 5 commits July 10, 2025 16:05

ruff/security fixes

e6682ac

Merge branch 'main' into add/change-diffusion-to-plugin

0e61123

bug fix

b6d4ba3

fix for job FAILs

77e6391

Merge branch 'main' into add/change-diffusion-to-plugin

eb489a4

Sourenm marked this pull request as ready for review July 11, 2025 17:31

Sourenm added 2 commits July 11, 2025 11:21

Merge branch 'main' into add/change-diffusion-to-plugin

a5b7198

Merge branch 'main' into add/change-diffusion-to-plugin

bb399dc

deep1401 requested changes Jul 11, 2025

View reviewed changes

Sourena and others added 4 commits July 13, 2025 13:36

change requests

b9fa272

Merge branch 'main' into add/change-diffusion-to-plugin

7ffb624

Merge branch 'main' into add/change-diffusion-to-plugin

53e42a4

Merge branch 'main' into add/change-diffusion-to-plugin

fefb5d9

deep1401 requested changes Jul 14, 2025

View reviewed changes

deep1401 and others added 2 commits July 15, 2025 10:09

Merge branch 'main' of https://github.com/transformerlab/transformerl…

35b80de

…ab-api into add/change-diffusion-to-plugin

Change requests

094d0dd

github-advanced-security bot found potential problems Jul 15, 2025

View reviewed changes

transformerlab/routers/experiment/diffusion.py Fixed Show fixed Hide fixed

Sourena and others added 6 commits July 15, 2025 10:43

security fix

475c55b

Merge branch 'main' into add/change-diffusion-to-plugin

3bc780d

remove diffusion_worker.py from shared

7ac521c

fix for correctly getting generation_id in case it's not in request body

2f0cf3c

Merge branch 'main' into add/change-diffusion-to-plugin

864d7dc

Merge branch 'main' into add/change-diffusion-to-plugin

654541a

deep1401 requested changes Jul 16, 2025

View reviewed changes

Sourena and others added 3 commits July 16, 2025 07:55

check for plugin installation, remove logger class

1b30620

Merge branch 'main' into add/change-diffusion-to-plugin

a8fb8e0

Merge branch 'main' into add/change-diffusion-to-plugin

bc1ff7a

deep1401 self-requested a review July 21, 2025 14:40

Merge branch 'main' into add/change-diffusion-to-plugin

8a613a0


		router = APIRouter(prefix="/diffusion", tags=["diffusion"])

		UNIVERSIAL_GENERATION_ID = str(uuid.uuid4())

Diffusion Plugin #421

Are you sure you want to change the base?

Diffusion Plugin #421

Uh oh!

Conversation

Sourenm commented Jul 8, 2025

Uh oh!

Uh oh!

Uh oh!

codecov bot commented Jul 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

codecov bot commented Jul 8, 2025 •

edited

Loading