建可迭代去噪循環(huán)的循環(huán)包裝器)
深入 Diffusers 模塊化管線之 LoopSequentialPipelineBlocks構(gòu)建可迭代去噪循環(huán)的循環(huán)包裝器【免費下載鏈接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.項目地址: https://gitcode.com/GitHub_Trending/di/diffusersLoopSequentialPipelineBlocks是 Diffusers 模塊化管線Modular Pipelines體系中的核心多塊類型之一它把一組 [~modular_pipelines.ModularPipelineBlocks] 組織成循環(huán)結(jié)構(gòu)讓數(shù)據(jù)通過intermediate_inputs/intermediate_outputs在每次迭代中循環(huán)流動是構(gòu)建默認就是迭代式去噪循環(huán)這類管線的標準手段。本指南將帶你從零創(chuàng)建一個循環(huán)包裝器與循環(huán)塊并用from_blocks_dict組裝出可運行的LoopSequentialPipelineBlocks最后結(jié)合倉庫中 Wan 視頻模型的真實去噪循環(huán)源碼理解循環(huán)結(jié)構(gòu)在生成管線中的實際落地方式。從多塊類型說起順序 vs 循環(huán)在模塊化管線的世界觀里管道是由一個個 [~modular_pipelines.ModularPipelineBlocks] 拼裝而成的而多塊類型multi-block types負責(zé)定義這些塊之間的編排方式SequentialPipelineBlocks將多個塊按順序串接數(shù)據(jù)線性地從上一個塊流向下一個塊LoopSequentialPipelineBlocks將多個塊以循環(huán)方式組合數(shù)據(jù)在塊之間循環(huán)流動每個塊被迭代執(zhí)行多次此外還有ConditionalPipelineBlocks/AutoPipelineBlocks這類根據(jù)輸入條件選擇執(zhí)行分支的多塊類型。LoopSequentialPipelineBlocks之所以重要是因為擴散模型的去噪過程本身就是一次迭代同一組預(yù)測噪聲 → 更新 latent的操作要在num_inference_steps個時間步上反復(fù)執(zhí)行。如果只用順序塊你就不得不把去噪循環(huán)手寫進某個單一塊內(nèi)部而循環(huán)塊則把循環(huán)體拆成多個可獨立替換的子塊循環(huán)邏輯與每輪具體做什么被徹底解耦。本文配套的模塊化管線總覽與快速上手可以幫你先建立整體認知本文聚焦LoopSequentialPipelineBlocks這一個點。循環(huán)包裝器定義循環(huán)結(jié)構(gòu)與迭代邏輯[~modular_pipelines.LoopSequentialPipelineBlocks] 也被稱為循環(huán)包裝器loop wrapper因為它負責(zé)定義三件事循環(huán)結(jié)構(gòu)、迭代變量和配置。你不需要像寫普通塊那樣覆寫inputs/intermediate_outputs屬性而是覆寫一套以loop_前綴命名的等價屬性。循環(huán)包裝器內(nèi)需要定義的變量如下表所示變量作用對應(yīng)普通塊的屬性loop_inputs用戶提供的值例如迭代步數(shù)num_steps[~modular_pipelines.ModularPipelineBlocks.inputs]loop_intermediate_inputs來自 [~modular_pipelines.PipelineState] 的中間變量供循環(huán)體消費[~modular_pipelines.ModularPipelineBlocks.intermediate_inputs]loop_intermediate_outputs由循環(huán)內(nèi)塊創(chuàng)建、最終寫回 [~modular_pipelines.PipelineState] 的新中間變量[~modular_pipelines.ModularPipelineBlocks.intermediate_outputs]__call__定義循環(huán)結(jié)構(gòu)和迭代邏輯是整個包裝器的核心——下面的示例定義了一個名為LoopWrapper的最小循環(huán)包裝器它聲明了唯一的循環(huán)輸入num_steps并在__call__中用一個for循環(huán)迭代每次迭代調(diào)用self.loop_step(...)按順序執(zhí)行所有注冊的子塊import torch from diffusers.modular_pipelines import LoopSequentialPipelineBlocks, ModularPipelineBlocks, InputParam, OutputParam class LoopWrapper(LoopSequentialPipelineBlocks): model_name test property def description(self): return Im a loop!! property def loop_inputs(self): return [InputParam(namenum_steps)] torch.no_grad() def __call__(self, components, state): block_state self.get_block_state(state) # 循環(huán)結(jié)構(gòu) - 可以根據(jù)您的需求定制 for i in range(block_state.num_steps): # loop_step 按順序執(zhí)行所有注冊的塊 components, block_state self.loop_step(components, block_state, ii) self.set_block_state(state, block_state) return components, state注意幾個關(guān)鍵細節(jié)loop_step與set_block_state的分工loop_step內(nèi)部已經(jīng)遍歷并執(zhí)行了sub_blocks中的全部子塊但整個循環(huán)結(jié)束后包裝器還要顯式調(diào)用self.set_block_state(state, block_state)把循環(huán)體內(nèi)累積的BlockState寫回全局的PipelineState這是循環(huán)包裝器與普通塊__call__結(jié)構(gòu)上最大的不同。迭代參數(shù)可以透傳循環(huán)包裝器可以把額外的參數(shù)如當(dāng)前迭代索引i通過loop_step(components, block_state, ii)傳給循環(huán)塊。從源碼可以看到loop_step的實現(xiàn)是依次調(diào)用每個子塊并把**kwargs原樣透傳def loop_step(self, components, state: PipelineState, **kwargs): for block_name, block in self.sub_blocks.items(): try: components, state block(components, state, **kwargs) except Exception as e: ... return components, state因此循環(huán)塊只要在__call__簽名里聲明i: int甚至t: torch.Tensor就能拿到每次迭代的上下文。循環(huán)塊每次迭代執(zhí)行的葉子單元循環(huán)塊本質(zhì)上仍然是一個 [~modular_pipelines.ModularPipelineBlocks]但與普通塊相比它的__call__方法行為有三個顯著區(qū)別它從循環(huán)包裝器接收調(diào)用而不是從順序容器接收調(diào)用它直接與 [~modular_pipelines.BlockState] 協(xié)作而不是與 [~modular_pipelines.PipelineState] 協(xié)作它不需要自行檢索或更新 [~modular_pipelines.BlockState] —— 該狀態(tài)由循環(huán)包裝器統(tǒng)一管理。循環(huán)塊共享同一個 [~modular_pipelines.BlockState]這正是值可以在循環(huán)的每次迭代中累積和變化的機制來源WanLoopBeforeDenoiser準備本輪輸入、WanLoopDenoiser算出noise_pred、WanLoopAfterDenoiser更新latents同一個BlockState在迭代間被反復(fù)讀寫。class LoopBlock(ModularPipelineBlocks): model_name test property def inputs(self): return [InputParam(namex)] property def intermediate_outputs(self): # 這個塊產(chǎn)生的輸出 return [OutputParam(namex)] property def description(self): return 我是一個在LoopWrapper類內(nèi)部使用的塊 def __call__(self, components, block_state, i: int): block_state.x 1 return components, block_state上面的LoopBlock每被調(diào)用一次就把block_state.x加一。注意它的__call__簽名是(self, components, block_state, i: int)接收的是block_state而不是state且直接返回(components, block_state)——這正是它與普通塊的__call__結(jié)構(gòu)檢索BlockState→ 計算 → 寫回PipelineState的差異所在。組裝用 from_blocks_dict 注冊循環(huán)塊要得到一個可運行的 [~modular_pipelines.LoopSequentialPipelineBlocks]使用類方法 [~modular_pipelines.LoopSequentialPipelineBlocks.from_blocks_dict] 把循環(huán)塊注冊進循環(huán)包裝器即可。blocks_dict的鍵是塊名值可以是塊類會自動實例化也可以是塊實例loop LoopWrapper.from_blocks_dict({block1: LoopBlock})如果想要在每次迭代中運行多個塊往字典里繼續(xù)添加即可。這允許你在完全不改動循環(huán)邏輯本身的前提下自由增刪、替換循環(huán)體內(nèi)的步驟loop LoopWrapper.from_blocks_dict({block1: LoopBlock(), block2: LoopBlock})從源碼實現(xiàn)看from_blocks_dict會先實例化傳入的類再統(tǒng)一設(shè)置block_classes、block_names與sub_blocksclassmethod def from_blocks_dict(cls, blocks_dict: dict[str, Any]) - LoopSequentialPipelineBlocks: instance cls() sub_blocks InsertableDict() for name, block in blocks_dict.items(): if inspect.isclass(block): sub_blocks[name] block() else: sub_blocks[name] block instance.block_classes [block.__class__ for block in blocks_dict.values()] instance.block_names list(blocks_dict.keys()) instance.sub_blocks blocks_dict return instance一個值得注意的約束是LoopSequentialPipelineBlocks.__init__會校驗所有子塊必須是葉子塊leaf blocks即不允許子塊自身再嵌套sub_blocks否則直接拋出ValueError。這是因為循環(huán)體內(nèi)的每一步應(yīng)當(dāng)是原子操作嵌套容器會讓循環(huán)語義變得不可控。因此循環(huán)塊只能是普通的ModularPipelineBlocks不能是SequentialPipelineBlocks或另一個ConditionalPipelineBlocks。源碼視角LoopSequentialPipelineBlocks 的完整契約回到src/diffusers/modular_pipelines/modular_pipeline.py中LoopSequentialPipelineBlocks類的實現(xiàn)可以提煉出它區(qū)別于SequentialPipelineBlocks的完整行為契約輸入合并_get_inputs()先并入loop_inputs再按順序并入各子塊的輸入跳過那些已被中間輸出覆蓋的變量required_inputs則取首個子塊必填 ∪ 循環(huán)必填 ∪ 其余子塊必填的并集。中間輸出合并intermediate_outputs先合并所有子塊的輸出再補充loop_intermediate_outputs中未出現(xiàn)的新變量outputs直接取最后一個子塊的intermediate_outputs。組件與配置透傳expected_components/expected_configs在聚合全部子塊的聲明之外還會追加循環(huán)包裝器自己聲明的loop_expected_components/loop_expected_configs——這使得調(diào)度器、guider 等循環(huán)級組件可以掛在包裝器上而不是重復(fù)聲明在每個子塊里。__call__必須由子類實現(xiàn)基類只拋出NotImplementedError__call__method needs to be implemented by the subclass也就是說循環(huán)的具體迭代方式如何遍歷 timesteps、是否顯示進度條、warmup 步數(shù)如何計算完全留給你的包裝器類定制。進度條工具基類提供了progress_bar(iterableNone, totalNone)與set_progress_bar_config(**kwargs)方便在長循環(huán)中渲染 tqdm 進度條并已用torch.compiler.disable標記避免與 torch.compile 沖突。真實案例Wan 視頻模型的迭代去噪循環(huán)倉庫中最能說明LoopSequentialPipelineBlocks實戰(zhàn)價值的案例是 Wan 系列視頻生成模型的去噪步驟位于src/diffusers/modular_pipelines/wan/denoise.py。首先定義一個循環(huán)包裝器WanDenoiseLoopWrapper它聲明了循環(huán)輸入timesteps與num_inference_steps并定制了循環(huán)邏輯遍歷每個 timestept把i和t同時透傳給循環(huán)體同時維護進度條與 warmup 步數(shù)class WanDenoiseLoopWrapper(LoopSequentialPipelineBlocks): model_name wan property def loop_inputs(self): return [ InputParam(timesteps, requiredTrue, type_hinttorch.Tensor, description...), InputParam(num_inference_steps, requiredTrue, type_hintint, description...), ] torch.no_grad() def __call__(self, components, state): block_state self.get_block_state(state) block_state.num_warmup_steps max( len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 ) with self.progress_bar(totalblock_state.num_inference_steps) as progress_bar: for i, t in enumerate(block_state.timesteps): components, block_state self.loop_step(components, block_state, ii, tt) if i len(block_state.timesteps) - 1 or ( (i 1) block_state.num_warmup_steps and (i 1) % components.scheduler.order 0 ): progress_bar.update() self.set_block_state(state, block_state) return components, state然后在包裝器的子類里通過block_classesblock_names聲明循環(huán)體的三個步驟把每輪迭代做什么配置化class WanDenoiseStep(WanDenoiseLoopWrapper): block_classes [ WanLoopBeforeDenoiser, # 準備 latent_model_input WanLoopDenoiser( # 帶 guidance 調(diào)用 transformer 預(yù)測噪聲 guider_input_fields{ encoder_hidden_states: (prompt_embeds, negative_prompt_embeds), } ), WanLoopAfterDenoiser, # 用 scheduler.step 更新 latents ] block_names [before_denoiser, denoiser, after_denoiser]三個子塊正是葉子塊 共享BlockState的典范WanLoopBeforeDenoiser.__call__(self, components, block_state, i, t)計算latent_model_inputWanLoopDenoiser通過 guider 拆分條件/無條件 batch 后調(diào)用transformer得到noise_predWanLoopAfterDenoiser調(diào)用components.scheduler.step(...)更新latents。同一個block_state.latents在數(shù)百次迭代中被反復(fù)讀出、更新、寫回最終得到去噪完成的 latent。Wan2.2 版本W(wǎng)an22DenoiseStep只需替換其中的 denoiser 子塊就能在同一循環(huán)邏輯下切換到雙 transformer / 雙 guider 的高噪聲與低噪聲兩階段去噪loop_expected_configs中的boundary_ratio即用于切分這兩個階段。除 Wan 之外倉庫中 FLUX、FLUX2、SDXL、SD3、LTX、HunyuanVideo1.5、QwenImage、Krea2、Ideogram4、MiniMax 等多個模型的denoise.py也都采用了同樣的循環(huán)包裝器模式如src/diffusers/modular_pipelines/flux/denoise.py印證了循環(huán)包裝器定義迭代結(jié)構(gòu)、循環(huán)塊定義每步操作這一設(shè)計在圖像與視頻生成管線中的普適性。與順序塊的分工總結(jié)維度SequentialPipelineBlocksLoopSequentialPipelineBlocks數(shù)據(jù)流向線性一次執(zhí)行循環(huán)多次迭代執(zhí)行子塊狀態(tài)管理每個子塊自行g(shù)et/set_block_state包裝器統(tǒng)一管理共享BlockState子塊直接使用子塊__call__簽名(components, state)(components, block_state, i, ...)子塊約束無特殊約束必須是葉子塊不能再嵌套sub_blocks典型用途編碼 → 去噪 → 解碼 的單遍步驟逐 timestep 迭代的去噪循環(huán)對于大多數(shù)一次性通過的步驟如文本編碼、latent 準備、VAE 解碼選擇SequentialPipelineBlocks對于天然迭代的計算去噪、擴散反向過程則應(yīng)該使用LoopSequentialPipelineBlocks。兩者的組合使用方式——順序塊把循環(huán)包裝器當(dāng)作其中一步嵌入整條管線——正是模塊化 Diffusers 用統(tǒng)一而可組合的積木搭建完整生成工作流的核心理念。相關(guān)測試可參考tests/modular_pipelines/test_modular_pipelines_custom_blocks.py中對from_blocks_dict與循環(huán)塊的驗證用例?!久赓M下載鏈接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.項目地址: https://gitcode.com/GitHub_Trending/di/diffusers創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考