commit 43d6fb36f11fdc8fdcf1cadb2fc013924bfca48d Author: Jesse Rabek Date: Thu Nov 27 18:27:31 2025 +0100 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f76d043 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +.DS_Store +__pycache__/ +.env* +*.pyc \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..36b2bfc --- /dev/null +++ b/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 + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..af0819b --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + diff --git a/bokeh_diffusion/__init__.py b/bokeh_diffusion/__init__.py new file mode 100644 index 0000000..c3cffa9 --- /dev/null +++ b/bokeh_diffusion/__init__.py @@ -0,0 +1,5 @@ +from .adapter_flux import BokehFluxControlAdapter +from .utils import color_transfer_lab + +__all__ = ["BokehFluxControlAdapter", "color_transfer_lab"] + diff --git a/bokeh_diffusion/adapter_flux.py b/bokeh_diffusion/adapter_flux.py new file mode 100644 index 0000000..cfd4602 --- /dev/null +++ b/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 diff --git a/bokeh_diffusion/utils.py b/bokeh_diffusion/utils.py new file mode 100644 index 0000000..83a8b2d --- /dev/null +++ b/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) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..c4b14d3 --- /dev/null +++ b/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", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..5e04e83 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +diffusers>=0.32.2 +transformers>=4.53.0 +huggingface_hub +torch>=2.0.0 +opencv-python +