Initial commit

.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+__pycache__/
+.env*
+*.pyc

README.md

@@ -0,0 +1,94 @@
+# ComfyUI Bokeh Diffusion
+
+ComfyUI custom nodes for Bokeh Diffusion - Defocus Blur Control in Text-to-Image Diffusion Models.
+
+This package provides ComfyUI nodes to add bokeh (defocus blur) control to FLUX models.
+
+## Installation
+
+1. Copy this directory (`ComfyUI_BokehDiffusion`) to your ComfyUI `custom_nodes` folder:
+   ```
+   cp -r ComfyUI_BokehDiffusion /path/to/ComfyUI/custom_nodes/
+   ```
+
+2. Install dependencies (if not already installed):
+   ```bash
+   pip install diffusers transformers huggingface_hub
+   ```
+
+## Usage
+
+### Basic Workflow
+
+1. **Load FLUX Model** - Use ComfyUI's standard model loading nodes to load a FLUX model
+2. **Load Bokeh Adapter** - Use the `Load Bokeh Adapter` node to load the adapter from HuggingFace
+3. **Apply Bokeh Adapter** - Connect your FLUX model and adapter to `Apply Bokeh Adapter`
+4. **Set Bokeh Level** - Use `Bokeh Level` node to set the desired bokeh level (0-30)
+5. **Generate** - Use `Bokeh KSampler` instead of the standard KSampler for generation
+
+### Node Descriptions
+
+#### Load Bokeh Adapter
+- **adapter_name**: HuggingFace model ID (default: "atfortes/BokehDiffusion")
+- **bokeh_scale**: Scale factor for bokeh effect (default: 1.0)
+- **lora_scale**: Scale factor for LoRA components (default: 1.0)
+
+#### Apply Bokeh Adapter
+- **model**: FLUX model from ComfyUI's model loader
+- **adapter**: Adapter from `Load Bokeh Adapter` node
+
+#### Bokeh Level
+- **bokeh_level**: Bokeh level from 0 (sharp) to 30 (maximum blur)
+- **bokeh_pivot**: (Optional) Pivot bokeh level for grounded generation
+- **num_grounding_steps**: (Optional) Number of grounding steps for scene consistency (0 = disabled, 18-24 recommended)
+
+#### Bokeh KSampler
+- **model**: Model with bokeh adapter applied
+- **positive**: Positive conditioning
+- **negative**: Negative conditioning
+- **latent_image**: Initial latent image
+- **seed**: Random seed
+- **steps**: Number of inference steps
+- **cfg**: Guidance scale
+- **sampler_name**: Sampler name (e.g., "euler")
+- **scheduler**: Scheduler name (e.g., "normal")
+- **bokeh_control**: Bokeh control from `Bokeh Level` node
+- **true_cfg**: (Optional) True CFG scale for negative prompt
+
+### Example Workflow
+
+```
+[Load Checkpoint] -> [Apply Bokeh Adapter] -> [Bokeh KSampler]
+                              ^
+                              |
+                    [Load Bokeh Adapter]
+
+[CLIP Text Encode] -> [Bokeh KSampler]
+                              ^
+                              |
+                    [Bokeh Level] -> [Bokeh KSampler]
+```
+
+## Features
+
+- **Unbounded Generation**: Control bokeh level from 0 to 30
+- **Grounded Generation**: Scene-consistent bokeh transitions using grounding steps
+- **Standard ComfyUI Integration**: Works with standard ComfyUI nodes (model loading, VAE, CLIP, etc.)
+
+## Notes
+
+- This node requires a FLUX model (e.g., FLUX.1-dev)
+- The adapter modifies the model's attention processors, so the model must support this
+- For best results with grounded generation, use 18-24 grounding steps
+- Bokeh level is normalized from 0-30 range to 0-1 internally
+
+## License
+
+This project is licensed under [NTU S-Lab License 1.0](https://github.com/atfortes/BokehDiffusion/blob/main/LICENSE).
+
+## Credits
+
+Based on the work by:
+- Armando Fortes, Tianyi Wei, Shangchen Zhou, Xingang Pan
+- Original repository: https://github.com/atfortes/BokehDiffusion
+

__init__.py

@@ -0,0 +1,4 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+

bokeh_diffusion/__init__.py

@@ -0,0 +1,5 @@
+from .adapter_flux import BokehFluxControlAdapter
+from .utils import color_transfer_lab
+
+__all__ = ["BokehFluxControlAdapter", "color_transfer_lab"]
+

bokeh_diffusion/adapter_flux.py

@@ -0,0 +1,368 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import FluxTransformer2DModel
+from huggingface_hub import PyTorchModelHubMixin
+from diffusers.models.activations import get_activation
+
+
+class LoRALinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        rank: int = 16,
+        network_alpha: float = None,
+        device: torch.device = None,
+        dtype: torch.dtype = None
+    ):
+        super().__init__()
+
+        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
+        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
+        self.network_alpha = network_alpha
+        self.rank = rank
+
+        nn.init.normal_(self.down.weight, std=1 / rank)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, hidden_states):
+        orig_dtype = hidden_states.dtype
+        dtype = self.down.weight.dtype
+
+        down_hidden_states = self.down(hidden_states.to(dtype))
+        up_hidden_states = self.up(down_hidden_states)
+
+        if self.network_alpha is not None:
+            up_hidden_states *= self.network_alpha / self.rank
+
+        return up_hidden_states.to(orig_dtype)
+
+
+class FluxAttnProcessor2_0(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, bokeh_embeds=None, attention_mask=None, image_rotary_emb=None):
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+        
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+class BokehFluxAttnProcessor2_0(torch.nn.Module):
+    def __init__(
+        self,
+        context_dim,
+        hidden_dim,
+        block_name=None,
+        bokeh_scale=1.0,
+        unfreeze_q=False,
+        unfreeze_k=False,
+        lora_rank=None,
+        lora_scale=1.0,
+        lora_alpha=None
+    ):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("BokehFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.block_name = block_name
+        self.hidden_dim = hidden_dim
+        self.context_dim = context_dim
+        self.bokeh_scale = bokeh_scale
+        self.lora_scale = lora_scale
+        self.unfreeze_q = unfreeze_q
+        self.unfreeze_k = unfreeze_k
+
+        self.to_k_bokeh = LoRALinear(context_dim, hidden_dim, rank=lora_rank, network_alpha=lora_alpha)
+        self.to_v_bokeh = LoRALinear(context_dim, hidden_dim, rank=lora_rank, network_alpha=lora_alpha)
+        if self.unfreeze_q:
+            self.to_q_adp = LoRALinear(hidden_dim, hidden_dim, rank=lora_rank, network_alpha=lora_alpha)
+        if self.unfreeze_k:
+            self.to_k_adp = LoRALinear(hidden_dim, hidden_dim, rank=lora_rank, network_alpha=lora_alpha)
+
+    def __call__(
+        self,
+        attn,
+        hidden_states,
+        encoder_hidden_states=None,
+        bokeh_embeds=None,
+        attention_mask=None,
+        image_rotary_emb=None,
+        perform_swap=False, 
+        batch_swap_ids=None
+    ):
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)  # [batch_size, seq_len, hidden_dim]
+        key = attn.to_k(hidden_states)  # [batch_size, seq_len, hidden_dim]
+        value = attn.to_v(hidden_states)  # [batch_size, seq_len, hidden_dim]
+
+        if self.unfreeze_q:
+            query = query + self.lora_scale * self.to_q_adp(query)
+        if self.unfreeze_k:
+            key = key + self.lora_scale * self.to_k_adp(key)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+        
+        if perform_swap:
+            query = query[batch_swap_ids]
+
+        if perform_swap and self.unfreeze_k:
+            key_swaps  = key[batch_swap_ids]
+            key  = torch.cat([key, key_swaps], dim=2)
+            value = torch.cat([value, value], dim=2)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if bokeh_embeds is not None:
+            key_dof = self.to_k_bokeh(bokeh_embeds)
+            value_dof = self.to_v_bokeh(bokeh_embeds)
+            key_dof = key_dof.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            value_dof = value_dof.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            camera_hidden_states = F.scaled_dot_product_attention(query, key_dof, value_dof, attn_mask=None, dropout_p=0.0, is_causal=False)
+            camera_hidden_states = camera_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+            camera_hidden_states = camera_hidden_states.to(query.dtype)
+            hidden_states = hidden_states + self.bokeh_scale * camera_hidden_states
+        
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+class BokehFluxControlAdapter(
+    torch.nn.Module,
+    PyTorchModelHubMixin,
+):
+    def __init__(
+        self,
+        base_model,
+        blocks,
+        attn_hidden_dim=3072,
+        bokeh_input_dim=1,
+        bokeh_hidden_dim=64,
+        bokeh_context_dim=768,
+        lora_rank=128,
+        lora_alpha=None,
+        lora_scale=1.0,
+        bokeh_scale=1.0,
+        unfreeze_q=False,
+        unfreeze_k=False,
+    ):
+        super().__init__()
+        self.blocks = blocks
+        self.bokeh_scale = bokeh_scale
+        self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha
+        self.lora_scale = lora_scale
+        self.attn_hidden_dim = attn_hidden_dim
+        self.bokeh_input_dim = bokeh_input_dim
+        self.bokeh_hidden_dim = bokeh_hidden_dim
+        self.bokeh_context_dim = bokeh_context_dim
+        self.unfreeze_q = unfreeze_q
+        self.unfreeze_k = unfreeze_k
+
+        self.embedding_layer = nn.Sequential(
+            nn.Linear(self.bokeh_input_dim, self.bokeh_hidden_dim),
+            get_activation("relu"),
+            nn.Linear(self.bokeh_hidden_dim, self.bokeh_context_dim)
+        )
+                
+        self._attach_transformer(base_model, blocks)
+
+    def _attach_transformer(self, base_model, blocks):
+        attn_procs = {}
+        for name in base_model.attn_processors.keys():
+            block_name = name.split(".attn")[0]
+            if block_name in blocks:
+                attn_procs[name] = BokehFluxAttnProcessor2_0(
+                    context_dim=self.bokeh_context_dim,
+                    hidden_dim=self.attn_hidden_dim,
+                    block_name=block_name,
+                    lora_rank=self.lora_rank,
+                    lora_alpha=self.lora_alpha,
+                    unfreeze_q=self.unfreeze_q,
+                    unfreeze_k=self.unfreeze_k,
+                )
+            else:
+                attn_procs[name] = FluxAttnProcessor2_0()
+        
+        base_model.set_attn_processor(attn_procs)
+        self.adapter_modules = torch.nn.ModuleList(base_model.attn_processors.values())
+
+    def set_lora_scale(self, lora_scale):
+        self.lora_scale = lora_scale
+        for module in self.adapter_modules:
+            module.lora_scale = lora_scale
+
+    def set_bokeh_scale(self, bokeh_scale):
+        self.bokeh_scale = bokeh_scale
+        for module in self.adapter_modules:
+            module.bokeh_scale = bokeh_scale
+
+    def forward(
+        self,
+        base_model: FluxTransformer2DModel,
+        bokeh_ann: torch.Tensor,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor,
+        pooled_projections: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        txt_ids: torch.Tensor,
+        img_ids: torch.Tensor,
+        perform_swap: bool = False,
+        batch_swap_ids: list[int] = None
+    ):
+        bokeh_embeds = self.embedding_layer(bokeh_ann)
+        noise_pred = base_model(
+            hidden_states,
+            timestep=timestep,
+            guidance=guidance,
+            pooled_projections=pooled_projections,
+            encoder_hidden_states=encoder_hidden_states,
+            txt_ids=txt_ids,
+            img_ids=img_ids,
+            joint_attention_kwargs={
+                "bokeh_embeds": bokeh_embeds,
+                "perform_swap": perform_swap,
+                "batch_swap_ids": batch_swap_ids,
+            },
+            return_dict=False,
+        )[0]
+        return noise_pred

bokeh_diffusion/utils.py

@@ -0,0 +1,23 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def color_transfer_lab(source_pil, target_pil, adjust_std=True):
+    source_np = np.array(source_pil, dtype=np.uint8)
+    target_np = np.array(target_pil, dtype=np.uint8)
+    source_lab = cv2.cvtColor(source_np, cv2.COLOR_RGB2LAB).astype(np.float32)
+    target_lab = cv2.cvtColor(target_np, cv2.COLOR_RGB2LAB).astype(np.float32)
+    mean_src, std_src = cv2.meanStdDev(source_lab)
+    mean_tgt, std_tgt = cv2.meanStdDev(target_lab)
+    eps = 1e-6
+    for c in range(3):
+        t_chan = target_lab[..., c]
+        t_chan -= mean_tgt[c][0]
+        if adjust_std and std_tgt[c][0] > eps:
+            t_chan *= (std_src[c][0] / std_tgt[c][0])
+        t_chan += mean_src[c][0]
+        target_lab[..., c] = t_chan
+    target_lab = np.clip(target_lab, 0, 255).astype(np.uint8)
+    matched = cv2.cvtColor(target_lab, cv2.COLOR_LAB2RGB)
+    return Image.fromarray(matched)

nodes.py

@@ -0,0 +1,381 @@
+import torch
+import numpy as np
+from typing import Optional, Tuple
+import sys
+import os
+
+# Import bokeh adapter
+from .bokeh_diffusion.adapter_flux import BokehFluxControlAdapter
+from .bokeh_diffusion.utils import color_transfer_lab
+
+MIN_BOKEH, MAX_BOKEH = 0, 30
+
+
+def find_flux_transformer(model):
+    """Find the FLUX transformer in a ComfyUI model"""
+    # Try different possible locations for the transformer
+    if hasattr(model, 'model'):
+        if hasattr(model.model, 'transformer'):
+            return model.model.transformer
+        elif hasattr(model.model, 'diffusion_model'):
+            # Check if it's a FluxTransformer2DModel
+            transformer = model.model.diffusion_model
+            if 'FluxTransformer2DModel' in str(type(transformer)) or hasattr(transformer, 'attn_processors'):
+                return transformer
+    elif hasattr(model, 'diffusion_model'):
+        transformer = model.diffusion_model
+        if 'FluxTransformer2DModel' in str(type(transformer)) or hasattr(transformer, 'attn_processors'):
+            return transformer
+    elif hasattr(model, 'transformer'):
+        return model.transformer
+
+    # Last resort: check if model itself is the transformer
+    if hasattr(model, 'attn_processors'):
+        return model
+
+    return None
+
+
+class BokehAdapterLoader:
+    """Load the Bokeh Diffusion adapter from HuggingFace"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "adapter_name": ("STRING", {
+                    "default": "atfortes/BokehDiffusion",
+                    "multiline": False
+                }),
+            },
+            "optional": {
+                "bokeh_scale": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 2.0,
+                    "step": 0.1
+                }),
+                "lora_scale": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 2.0,
+                    "step": 0.1
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("BOKEH_ADAPTER",)
+    RETURN_NAMES = ("adapter",)
+    FUNCTION = "load_adapter"
+    CATEGORY = "bokeh_diffusion"
+
+    def load_adapter(self, adapter_name: str, bokeh_scale: float = 1.0, lora_scale: float = 1.0):
+        """Load the bokeh adapter configuration"""
+        adapter_info = {
+            "adapter_name": adapter_name,
+            "bokeh_scale": bokeh_scale,
+            "lora_scale": lora_scale,
+            "adapter": None  # Will be loaded when model is available
+        }
+        return (adapter_info,)
+
+
+class ApplyBokehAdapter:
+    """Apply the Bokeh adapter to a FLUX model"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "adapter": ("BOKEH_ADAPTER",),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "apply_adapter"
+    CATEGORY = "bokeh_diffusion"
+
+    def apply_adapter(self, model, adapter):
+        """Apply the bokeh adapter to the model's transformer"""
+        try:
+            # Find the transformer
+            transformer = find_flux_transformer(model)
+
+            if transformer is None:
+                raise Exception("Could not find FLUX transformer in model. Make sure you're using a FLUX model.")
+
+            # Check if transformer has attn_processors (required for adapter)
+            if not hasattr(transformer, 'attn_processors'):
+                raise Exception("Model transformer does not have attn_processors. This may not be a FLUX model.")
+
+            # Load the adapter if not already loaded
+            adapter_name = adapter["adapter_name"]
+            if adapter["adapter"] is None:
+                print(f"Loading Bokeh adapter from {adapter_name}...")
+                adapter_obj = BokehFluxControlAdapter.from_pretrained(
+                    adapter_name,
+                    base_model=transformer,
+                )
+                # Move to same device as transformer
+                device = next(transformer.parameters()).device
+                dtype = next(transformer.parameters()).dtype
+                adapter_obj = adapter_obj.to(device=device, dtype=dtype)
+                adapter["adapter"] = adapter_obj
+                print("Bokeh adapter loaded successfully.")
+
+            adapter_obj = adapter["adapter"]
+
+            # Set scales
+            if adapter["bokeh_scale"] != 1.0:
+                adapter_obj.set_bokeh_scale(adapter["bokeh_scale"])
+            if adapter["lora_scale"] != 1.0:
+                adapter_obj.set_lora_scale(adapter["lora_scale"])
+
+            # Store adapter in model for use during sampling
+            model.bokeh_adapter = adapter_obj
+            model.bokeh_adapter_config = adapter
+            model._bokeh_transformer = transformer
+
+            return (model,)
+        except Exception as e:
+            raise Exception(f"Failed to apply bokeh adapter: {e}")
+
+
+class BokehLevel:
+    """Set bokeh level for generation"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "bokeh_level": ("FLOAT", {
+                    "default": 15.0,
+                    "min": MIN_BOKEH,
+                    "max": MAX_BOKEH,
+                    "step": 0.1,
+                    "tooltip": "Bokeh level from 0 (sharp) to 30 (maximum blur)"
+                }),
+            },
+            "optional": {
+                "bokeh_pivot": ("FLOAT", {
+                    "default": 15.0,
+                    "min": MIN_BOKEH,
+                    "max": MAX_BOKEH,
+                    "step": 0.1,
+                    "tooltip": "Pivot bokeh level for grounded generation (scene consistency)"
+                }),
+                "num_grounding_steps": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 50,
+                    "tooltip": "Number of grounding steps for scene consistency (0 = disabled, 18-24 recommended)"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("BOKEH_CONTROL",)
+    RETURN_NAMES = ("bokeh_control",)
+    FUNCTION = "create_control"
+    CATEGORY = "bokeh_diffusion"
+
+    def create_control(self, bokeh_level: float, bokeh_pivot: Optional[float] = None, num_grounding_steps: int = 0):
+        """Create bokeh control parameters"""
+        # Normalize bokeh level (0-30 -> 0-1)
+        bokeh_target = bokeh_level / MAX_BOKEH
+        bokeh_pivot_norm = bokeh_pivot / MAX_BOKEH if bokeh_pivot is not None else None
+
+        control = {
+            "bokeh_target": bokeh_target,
+            "bokeh_pivot": bokeh_pivot_norm,
+            "num_grounding_steps": num_grounding_steps,
+            "is_grounded": num_grounding_steps > 0 and bokeh_pivot is not None
+        }
+        return (control,)
+
+
+class BokehKSampler:
+    """KSampler with bokeh control support for FLUX models"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "latent_image": ("LATENT",),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 30, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.1}),
+                "sampler_name": ("STRING", {"default": "euler"}),
+                "scheduler": ("STRING", {"default": "normal"}),
+                "bokeh_control": ("BOKEH_CONTROL",),
+            },
+            "optional": {
+                "true_cfg": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 10.0,
+                    "step": 0.1,
+                    "tooltip": "True CFG scale for negative prompt"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("latent",)
+    FUNCTION = "sample"
+    CATEGORY = "bokeh_diffusion/sampling"
+
+    def sample(self, model, positive, negative, latent_image, seed, steps, cfg, sampler_name, scheduler,
+               bokeh_control, true_cfg=1.0):
+        """Sample with bokeh control"""
+
+        # Check if model has bokeh adapter
+        if not hasattr(model, 'bokeh_adapter'):
+            raise Exception("Model does not have bokeh adapter applied. Use ApplyBokehAdapter node first.")
+
+        bokeh_adapter = model.bokeh_adapter
+        transformer = model._bokeh_transformer
+
+        # Get bokeh control parameters
+        bokeh_target = bokeh_control["bokeh_target"]
+        bokeh_pivot = bokeh_control.get("bokeh_pivot")
+        num_grounding_steps = bokeh_control.get("num_grounding_steps", 0)
+        is_grounded = bokeh_control.get("is_grounded", False)
+
+        # Get latent
+        samples = latent_image["samples"]
+        device = samples.device
+        dtype = samples.dtype
+
+        # Prepare bokeh annotations
+        if is_grounded:
+            bokeh_ann = torch.tensor([bokeh_target, bokeh_pivot], dtype=dtype, device=device)
+        else:
+            bokeh_ann = torch.tensor([bokeh_target], dtype=dtype, device=device)
+        bokeh_ann = bokeh_ann.unsqueeze(1)
+        neg_bokeh_ann = torch.full_like(bokeh_ann, -1)
+
+        # Store bokeh info in transformer for use during forward pass
+        # The attention processors will access this via joint_attention_kwargs
+        transformer._bokeh_ann = bokeh_ann
+        transformer._neg_bokeh_ann = neg_bokeh_ann
+        transformer._is_grounded = is_grounded
+        transformer._num_grounding_steps = num_grounding_steps
+        transformer._true_cfg = true_cfg
+        transformer._current_step = 0
+
+        # Patch the transformer's forward method to inject bokeh annotations
+        original_forward = transformer.forward
+
+        def patched_forward(
+            hidden_states,
+            timestep=None,
+            encoder_hidden_states=None,
+            pooled_projections=None,
+            added_cond_kwargs=None,
+            return_dict=True,
+            **kwargs
+        ):
+            # Get current step info
+            current_step = getattr(transformer, '_current_step', 0)
+            is_grounded = getattr(transformer, '_is_grounded', False)
+            num_grounding_steps = getattr(transformer, '_num_grounding_steps', 0)
+            bokeh_ann = getattr(transformer, '_bokeh_ann', None)
+
+            # Prepare joint_attention_kwargs with bokeh info
+            joint_attention_kwargs = kwargs.get('joint_attention_kwargs', {})
+            if bokeh_ann is not None:
+                perform_swap = is_grounded and (current_step < num_grounding_steps)
+                batch_swap_ids = [1, 1] if is_grounded else None
+
+                # Get bokeh embeddings from adapter
+                bokeh_embeds = bokeh_adapter.embedding_layer(bokeh_ann)
+                joint_attention_kwargs['bokeh_embeds'] = bokeh_embeds
+                joint_attention_kwargs['perform_swap'] = perform_swap
+                joint_attention_kwargs['batch_swap_ids'] = batch_swap_ids
+
+            kwargs['joint_attention_kwargs'] = joint_attention_kwargs
+            return original_forward(
+                hidden_states,
+                timestep=timestep,
+                encoder_hidden_states=encoder_hidden_states,
+                pooled_projections=pooled_projections,
+                added_cond_kwargs=added_cond_kwargs,
+                return_dict=return_dict,
+                **kwargs
+            )
+
+        transformer.forward = patched_forward
+
+        try:
+            # Use standard ComfyUI sampling
+            import comfy.samplers
+            import comfy.sample
+
+            # Get the sampler function
+            sampler = comfy.samplers.sampler_object(sampler_name)
+
+            # Create a callback to track steps for grounded generation
+            def step_callback(step, x0, x, total_steps):
+                transformer._current_step = step
+
+            # Sample using ComfyUI's sample function
+            # Note: This is a simplified version. Full implementation may need
+            # to handle negative prompts with bokeh control separately
+            disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+
+            sampled = comfy.sample.sample(
+                model,
+                seed,
+                steps,
+                cfg,
+                sampler_name,
+                scheduler,
+                positive,
+                negative,
+                samples,
+                denoise=1.0,
+                disable_noise=False,
+                start_step=None,
+                last_step=None,
+                force_full_denoise=False,
+                callback=step_callback,
+                disable_pbar=disable_pbar,
+                seed_mode="scale_alike"
+            )
+
+            # Restore original forward
+            transformer.forward = original_forward
+
+            return (sampled,)
+
+        except ImportError:
+            # Fallback: if ComfyUI imports fail, return original samples
+            # This allows the node to be defined even if not in ComfyUI
+            print("Warning: ComfyUI not found. Bokeh control requires ComfyUI environment.")
+            transformer.forward = original_forward
+            return ({"samples": samples},)
+        except Exception as e:
+            # Restore original forward on error
+            transformer.forward = original_forward
+            raise Exception(f"Sampling failed: {e}")
+
+
+# Node class mappings for ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "BokehAdapterLoader": BokehAdapterLoader,
+    "ApplyBokehAdapter": ApplyBokehAdapter,
+    "BokehLevel": BokehLevel,
+    "BokehKSampler": BokehKSampler,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BokehAdapterLoader": "Load Bokeh Adapter",
+    "ApplyBokehAdapter": "Apply Bokeh Adapter",
+    "BokehLevel": "Bokeh Level",
+    "BokehKSampler": "Bokeh KSampler",
+}

requirements.txt

@@ -0,0 +1,6 @@
+diffusers>=0.32.2
+transformers>=4.53.0
+huggingface_hub
+torch>=2.0.0
+opencv-python
+