207 lines
8.0 KiB
Python
207 lines
8.0 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Application service for selectable tiled diffusion latent sampling."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, ClassVar
|
|
|
|
from ..domain.conditioning_batch import ConditioningBatch
|
|
from ..domain.noise_inversion import NoiseInversionOptions
|
|
from ..domain.regional_features import (
|
|
EMPTY_REGIONAL_FEATURE_REQUEST,
|
|
TILED_DIFFUSION_REGIONAL_SAMPLER_CAPABILITIES,
|
|
RegionalFeature,
|
|
RegionalFeatureRequest,
|
|
)
|
|
from ..domain.tiled_diffusion import validate_tiled_diffusion_mode
|
|
from ..runtime.detail_previews import DetailPreviewContext
|
|
from .regional_capability_admission_service import (
|
|
RegionalCapabilityAdmissionService,
|
|
)
|
|
from .regional_sampling_preparation_service import (
|
|
RegionalSamplingPreparationService,
|
|
)
|
|
from .regional_tiled_diffusion_sampling_service import (
|
|
RegionalTiledDiffusionSamplingService,
|
|
)
|
|
from .segs_guided_tiled_diffusion_sampling_service import (
|
|
SEGSGuidedTiledDiffusionSamplingService,
|
|
)
|
|
from .tiled_diffusion_conditioning_batch_service import (
|
|
TiledDiffusionConditioningBatchService,
|
|
)
|
|
from .tiled_diffusion_item_sampling_service import TiledDiffusionItemSamplingService
|
|
|
|
Latent = dict[str, Any]
|
|
|
|
|
|
class TiledDiffusionSamplingService:
|
|
"""Route tiled diffusion sampling requests to the selected runtime."""
|
|
|
|
regional_preparation_service_class: ClassVar[
|
|
type[RegionalSamplingPreparationService]
|
|
] = RegionalSamplingPreparationService
|
|
capability_admission_service_class: ClassVar[
|
|
type[RegionalCapabilityAdmissionService]
|
|
] = RegionalCapabilityAdmissionService
|
|
regional_sampling_service_class: ClassVar[
|
|
type[RegionalTiledDiffusionSamplingService]
|
|
] = RegionalTiledDiffusionSamplingService
|
|
item_sampling_service_class: ClassVar[type[TiledDiffusionItemSamplingService]] = (
|
|
TiledDiffusionItemSamplingService
|
|
)
|
|
segs_sampling_service_class: ClassVar[
|
|
type[SEGSGuidedTiledDiffusionSamplingService]
|
|
] = SEGSGuidedTiledDiffusionSamplingService
|
|
conditioning_batch_service_class: ClassVar[
|
|
type[TiledDiffusionConditioningBatchService]
|
|
] = TiledDiffusionConditioningBatchService
|
|
|
|
def sample(
|
|
self,
|
|
*,
|
|
diffusion_mode: str,
|
|
model: Any,
|
|
seed: int,
|
|
steps: int,
|
|
cfg: float,
|
|
sampler_name: str,
|
|
scheduler: str,
|
|
positive: Any,
|
|
negative: Any,
|
|
latent_image: Latent,
|
|
denoise: float,
|
|
latent_tile_width: int,
|
|
latent_tile_height: int,
|
|
latent_tile_overlap: int,
|
|
latent_tile_batch_size: int,
|
|
preview_context: DetailPreviewContext | None = None,
|
|
differential_diffusion: bool = False,
|
|
feature_request: RegionalFeatureRequest = EMPTY_REGIONAL_FEATURE_REQUEST,
|
|
segs: object | None = None,
|
|
region_masks: object | None = None,
|
|
regional_prompt_weight: float = 0.5,
|
|
region_mask_feather: int = 0,
|
|
noise_inversion: NoiseInversionOptions | None = None,
|
|
) -> Latent:
|
|
"""Sample a latent with the selected tiled diffusion method."""
|
|
|
|
validate_tiled_diffusion_mode(diffusion_mode)
|
|
regional = self.regional_preparation_service_class().prepare(
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_image,
|
|
region_masks=region_masks,
|
|
regional_prompt_weight=regional_prompt_weight,
|
|
region_mask_feather=region_mask_feather,
|
|
)
|
|
effective_request = feature_request
|
|
if regional.active:
|
|
effective_request = effective_request.with_feature(
|
|
RegionalFeature.FULL_CONTEXT_MASKED_CONDITIONING
|
|
)
|
|
capability_admission = self.capability_admission_service_class().admit(
|
|
request=effective_request,
|
|
sampler_capabilities=TILED_DIFFUSION_REGIONAL_SAMPLER_CAPABILITIES,
|
|
model=model,
|
|
)
|
|
if regional.mask_bank is not None:
|
|
return self.regional_sampling_service_class().sample(
|
|
item_sampler=self.item_sampling_service_class().sample,
|
|
region_masks=regional.mask_bank.planning_masks,
|
|
segs=segs,
|
|
diffusion_mode=diffusion_mode,
|
|
model=model,
|
|
seed=seed,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=regional.positive,
|
|
negative=regional.negative,
|
|
latent_image=latent_image,
|
|
denoise=denoise,
|
|
latent_tile_width=latent_tile_width,
|
|
latent_tile_height=latent_tile_height,
|
|
latent_tile_overlap=latent_tile_overlap,
|
|
latent_tile_batch_size=latent_tile_batch_size,
|
|
preview_context=preview_context,
|
|
differential_diffusion=differential_diffusion,
|
|
capability_admission=capability_admission,
|
|
noise_inversion=noise_inversion,
|
|
)
|
|
if segs is not None:
|
|
return self.segs_sampling_service_class().sample(
|
|
item_sampler=self.item_sampling_service_class().sample,
|
|
diffusion_mode=diffusion_mode,
|
|
model=model,
|
|
seed=seed,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_image,
|
|
denoise=denoise,
|
|
latent_tile_width=latent_tile_width,
|
|
latent_tile_height=latent_tile_height,
|
|
latent_tile_overlap=latent_tile_overlap,
|
|
latent_tile_batch_size=latent_tile_batch_size,
|
|
preview_context=preview_context,
|
|
differential_diffusion=differential_diffusion,
|
|
capability_admission=capability_admission,
|
|
noise_inversion=noise_inversion,
|
|
segs=segs,
|
|
)
|
|
if isinstance(positive, ConditioningBatch) or isinstance(
|
|
negative,
|
|
ConditioningBatch,
|
|
):
|
|
return self.conditioning_batch_service_class().sample(
|
|
item_sampler=self.item_sampling_service_class().sample,
|
|
diffusion_mode=diffusion_mode,
|
|
model=model,
|
|
seed=seed,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_image,
|
|
denoise=denoise,
|
|
latent_tile_width=latent_tile_width,
|
|
latent_tile_height=latent_tile_height,
|
|
latent_tile_overlap=latent_tile_overlap,
|
|
latent_tile_batch_size=latent_tile_batch_size,
|
|
preview_context=preview_context,
|
|
differential_diffusion=differential_diffusion,
|
|
capability_admission=capability_admission,
|
|
noise_inversion=noise_inversion,
|
|
)
|
|
return self.item_sampling_service_class().sample(
|
|
diffusion_mode=diffusion_mode,
|
|
model=model,
|
|
seed=seed,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_image,
|
|
denoise=denoise,
|
|
latent_tile_width=latent_tile_width,
|
|
latent_tile_height=latent_tile_height,
|
|
latent_tile_overlap=latent_tile_overlap,
|
|
latent_tile_batch_size=latent_tile_batch_size,
|
|
preview_context=preview_context,
|
|
differential_diffusion=differential_diffusion,
|
|
capability_admission=capability_admission,
|
|
noise_inversion=noise_inversion,
|
|
)
|