[LatentFlow] Simplifying Complex Pipelines with a Functional API

[tumurzakov]

Hi everyone,

I've been delving into numerous pipelines in the Diffusers repository, and I've noticed a recurring pattern among them. Names like StableDiffusionPipeline, StableDiffusionInpaintPipeline, StableDiffusionControlNetPipeline, and even StableDiffusionDeathStarPipeline all seem to share the same code, repeatedly copy-pasted throughout these files.

In my exploration of AnimateDiff, I've encountered intricate challenges, particularly with multiple nested loops: FreeInit > Denoising Loop > WalkPrompt > RegionPrompt. The complexity of these nested structures has left me somewhat overwhelmed.

To address this, I've attempted to create a functional-style API for the Stable Diffusion pipeline. Below is an example of what I've come up with:

LatentFlow

unet = Unet(
    unet=pipe.unet.to('cuda'),
    scheduler=pipe.scheduler,
    guidance_scale=10,
)

state = State()
video = \
    (Video(video=torch.randn((16,288,512,3))) \
        | Debug("Video to video") \
        | VideoLoad('src/tests_func/bunny.mp4', device='cuda', video_length=48) \
        | VideoShow(fps=16) \
        | VaeVideoEncode(vae=pipe.vae.to('cuda')) \
        | Invert(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, unet=pipe.unet, scheduler=pipe.scheduler) \
        | LatentShow(fps=16, vae=pipe.vae.to('cuda')) \
        > state("latent") \
        ) >> \
    (Latent(latent=torch.randn((1,4,16,36,64)), device='cuda') \
        - Debug("Text to video") \
        - Noise(scheduler=pipe.scheduler) \
        - LatentShow(fps=16, vae=pipe.vae.to('cuda')) \
        < state("latent") \
        ) >> \
    (Prompt("a cat walking down city street") \
        | CompelPromptEncode(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder.to('cuda')) \
        > state("embeddings") \
        ) >> \
    (state \
        | Schedule(scheduler=pipe.scheduler, num_inference_steps=10, strength=0.7) \
        > state('timesteps') \
        ) >> \
    (state['latent'] \
        - Debug("Adding noise")
        - AddNoise(scheduler=pipe.scheduler) \
        < state('latent') \
        ) >> \
    (state
        | Debug("Diffusion loop") \
        | Diffuse(callback=lambda timestep, state: state | unet)
        | VaeLatentDecode(vae=pipe.vae.to('cuda'))
    ) \
    | VideoShow(fps=16)

• | is pipe (like unix pipe)
• > write value in state (like in unix too)
• - bypass for pipe
• < bypass for write
• >> go to next statement

I'd love to hear your thoughts on this API. Additionally, if you have any ideas on how it could be further improved or refined, please feel free to share them.

Looking forward to your insights and feedback.

[a-r-r-o-w]

ComfyUI like node system but with code? I'm all for using this in my workflows ❤️ Very nice approach that provides far more control to power users without the need for copying large amounts of code to customize a pipeline to your needs (which is what most of us do I think). I have been wanting to work on something similar for a while too, but more related to different denosing backends (attend-and-excite, self-attention-guidance, sd-safe, etc.), something like:

pipe.set_denoise_backend(DenoisingBackend.SAG)

It would also make it easier to add newer features like FreeU or FreeInit or StyleAligned or whatever accessible across all pipelines/tasks because State() can handle it all behind the scenes. From the user end, it would be as simple as adding an additional line (as you demonstrate above, FreeU(s1, s2, b1, b2), or FreeInit(use_fast_sampling=True)).

I personally prefer the style of sklearn-like composite pipelines, but using the proposed syntax would make me feel 10x smarter instantly 😆 Great share!

[tumurzakov]

Thank you for your thoughts!

You mean this pipelines

cols_trans = ColumnTransformer([
    ('txt', TfidfVectorizer(), text_col),
    ('txt_kw', CountVectorizer(analyzer=no_analyzer), topic_col),
    ('ohe', OneHotEncoder(drop='first'), cat_cols)
])

One of the hardest problems was to make denoising loop. I made it in recursion. I have no experience with sklearn pipelines and don't know if loops could be implemented there.

[tumurzakov]

Hey, @a-r-r-o-w I just made ComfyUI node support for my framework. Now it is literally ComfyUI for code.

[sayakpaul]

Thank you for initiating this discussion.

One of the hardest problems was to make denoising loop. I made it in recursion.

One of the reasons why we don't want to introduce too many abstractions. We manage the copy mechanism by using "# Copied from ..." statements. The forward() should encapsulate the core logic while also catering to the features well established in the community.

Other popular features are enabled with set_XXX() or enable_XXX() methods.

Also note that we support pipeline chaining seamlessly: https://huggingface.co/docs/diffusers/main/en/using-diffusers/img2img#chained-image-to-image-pipelines.

We also expose more explicit control over the most important components inside a pipeline via a callback mechanism: https://huggingface.co/docs/diffusers/main/en/using-diffusers/callback.

Cc: @yiyixuxu @patrickvonplaten @DN6 @stevhliu

[tumurzakov]

I'm experimenting with my API and it seems promising to me. As @a-r-r-o-w said that it looks like ComfyUI for code

(state['latent']
        | LoraOn(loras={
            f"{models}/lora/details.safetensors": 0.8,
            }, pipe=pipe)
        | Loop(state['timesteps'], callback=lambda timestep_index, timestep: 
            state
            | Loop(TileGenerator(Tile(36,64,48), state['latent']), callback=lambda tile: 
                state 
                | cnet(timestep_index, timestep, latent=tile, controlnet_scale=[1.0]) 
                | unet(timestep_index, timestep, latent=tile)
                )
            )
        | LoraOff(pipe=pipe)
        | VaeLatentDecode(vae=pipe.vae.to('cuda'))
    ) \

I'll add several more loops for FreeInit and RegionPrompt i think it will look like

(state['latent']
        | LoraOn(loras={
            f"{models}/lora/details.safetensors": 0.8,
            }, pipe=pipe)
        | Loop(range(freeinit_repeat), callback=lambda freeinit_repeat_index: 
           state
           | Loop(state['timesteps'], callback=lambda timestep_index, timestep: 
              state
              | Loop(TileGenerator(Tile(36,64,48), state['latent']), callback=lambda tile: 
                  | Loop(RegionPromptGenerator(), callback=lambda region: 
                    state 
                    | cnet(timestep_index, timestep, latent=tile, controlnet_scale=[1.0]) 
                    | unet(timestep_index, timestep, latent=tile)
                    ) | Debug("RegionPrompt loop end")
                  ) | Debug("Tile loop end")
              ) | Debug("Denoise loop end")
            ) | Debug("FreeInit loop end")
        | LoraOff(pipe=pipe)
        | VaeLatentDecode(vae=pipe.vae.to('cuda'))
    ) \

[tumurzakov]

Actually it is not so important how it looks like. More important thing is that concepts must be some kind of blocks. If you look at any arxiv paper in diffusion theme with code release, they will copy-paste one of diffusers pipeline and add something new. Code often obsolete. For example MotionCtrl. They made their module for VideoCrafter and want to make it for AnimateDiff, but because every project releases their own pipelines it is hard to do. Better if the could release only some block that easily could be integrated in general pipeline.

[tumurzakov]

Hello everybody!

I've just added support for ComfyUI nodes in my LatentFlow framework (ComfyUI like nodes for code). While it's primarily designed for video processing, it can also generate images.

import sys

sys.path.insert(0, f'/path/to/latentflow/src')
sys.path.insert(0, f'/path/to/ComfyUI')

import latentflow as lf
import nodes
import torch

state = lf.State({})

(lf.ComfyNode(nodes.CheckpointLoaderSimple, ckpt_name='revanimated.safetensors').apply()
    | lf.Set(state, "checkpoint")
    | lf.ComfyNode(nodes.EmptyLatentImage, width=512, height=512, batch_size=1)
    | lf.ComfyNode(
        nodes.KSampler,
        model=state['checkpoint']['MODEL'],
        positive=lf.ComfyNode(nodes.CLIPTextEncode, text="a cat", clip=state["checkpoint"]["CLIP"]).apply(),
        negative=lf.ComfyNode(nodes.CLIPTextEncode, text="", clip=state["checkpoint"]["CLIP"]).apply(),
    )
    | lf.ComfyNode(nodes.VAEDecode, vae=state['checkpoint']["VAE"])
    | lf.Apply(lambda x: lf.Video('HWC', torch.stack([x['IMAGE']])))
    | lf.VideoShow()
)

• | is pipe (like unix pipe)
• > write value in state (like in unix too)
• - bypass for pipe
• < bypass for write
• >> go to next statement

I would greatly appreciate your feedback. While it's in the early stages and lacks documentation, I've been extensively working with it for my own purposes.

[a-r-r-o-w]

Hey, great work! This looks very promising and I'd love to try it out when I'm more free.

[sayakpaul]

Great work here @tumurzakov! I love how you have gone about the design. If there's a tweet or announcement, we would love to share it!