diff --git a/__init__.py b/__init__.py index ab64490..65c9e17 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,13 @@ from .get_image_size_ratio_node import GetImageSizeRatio from .noise_plus_blend_node import NoisePlusBlend from .integer_settings_node import IntegerSettings from .choose_upscale_model_node import ChooseUpscaleModel +from .region_mask_generator_node import RegionMaskGenerator +from .region_mask_validator_node import RegionMaskValidator +from .region_mask_processor_node import RegionMaskProcessor +from .region_mask_conditioning_node import RegionMaskConditioning +from .flux_attention_control_node import FluxAttentionControl +from .region_overlay_visualizer_node import RegionOverlayVisualizer +from .flux_attention_cleanup_node import FluxAttentionCleanup print("Initializing ControlAltAI Nodes") @@ -20,6 +27,13 @@ NODE_CLASS_MAPPINGS = { "NoisePlusBlend": NoisePlusBlend, "IntegerSettings": IntegerSettings, "ChooseUpscaleModel": ChooseUpscaleModel, + "RegionMaskGenerator": RegionMaskGenerator, + "RegionMaskValidator": RegionMaskValidator, + "RegionMaskProcessor": RegionMaskProcessor, + "RegionMaskConditioning": RegionMaskConditioning, + "FluxAttentionControl": FluxAttentionControl, + "RegionOverlayVisualizer": RegionOverlayVisualizer, + "FluxAttentionCleanup": FluxAttentionCleanup, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -32,4 +46,11 @@ NODE_DISPLAY_NAME_MAPPINGS = { "NoisePlusBlend": "Noise Plus Blend", "IntegerSettings": "Integer Settings", "ChooseUpscaleModel": "Choose Upscale Model", -} + "RegionMaskGenerator": "Region Mask Generator", + "RegionMaskValidator": "Region Mask Validator", + "RegionMaskProcessor": "Region Mask Processor", + "RegionMaskConditioning": "Region Mask Conditioning", + "FluxAttentionControl": "Flux Attention Control", + "RegionOverlayVisualizer": "Region Overlay Visualizer", + "FluxAttentionCleanup": "Flux Attention Cleanup", +} \ No newline at end of file diff --git a/flux_attention_cleanup_node.py b/flux_attention_cleanup_node.py new file mode 100644 index 0000000..ab7e49f --- /dev/null +++ b/flux_attention_cleanup_node.py @@ -0,0 +1,75 @@ +import torch +from comfy.ldm.modules import attention as comfy_attention +from comfy.ldm.flux import math as flux_math +from comfy.ldm.flux import layers as flux_layers + +class AnyType(str): + """A special class that is always equal in not equal comparisons""" + def __ne__(self, __value: object) -> bool: + return False + +any_type = AnyType("*") + +class FluxAttentionCleanup: + def __init__(self): + self.original_attention = comfy_attention.optimized_attention + self.original_flux_attention = flux_math.attention + self.original_flux_layers_attention = flux_layers.attention + self.current_attn_mask = None + print("FluxAttentionCleanup initialized") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "any_input": (any_type, {}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("message",) + FUNCTION = "cleanup_attention" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def cleanup_attention(self, any_input): + """Skip cleanup during normal operation, but clean on workflow switch""" + message = "Attention preserved for current workflow. Will clean on workflow switch." + print("\n" + message) + return (message,) + + def __del__(self): + """Clean up attention when switching workflows""" + try: + print("\nStarting attention cleanup for workflow switch...") + + # Reset attention functions to original state + flux_math.attention = self.original_flux_attention + flux_layers.attention = self.original_flux_layers_attention + + # Clear attention mask + if hasattr(flux_math.attention, 'keywords'): + if 'attn_mask' in flux_math.attention.keywords: + flux_math.attention.keywords['attn_mask'] = None + + # Clear stored mask + if self.current_attn_mask is not None: + del self.current_attn_mask + self.current_attn_mask = None + + # Force CUDA cleanup + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + print("Workflow switch: Region Attention Cleanup Successful") + except: + pass + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "FluxAttentionCleanup": FluxAttentionCleanup +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxAttentionCleanup": "Flux Attention Cleanup" +} \ No newline at end of file diff --git a/flux_attention_control_node.py b/flux_attention_control_node.py new file mode 100644 index 0000000..df3bf17 --- /dev/null +++ b/flux_attention_control_node.py @@ -0,0 +1,304 @@ +import torch +from torch import Tensor +import torch.nn.functional as F +from typing import List, Dict, Optional, Tuple +from einops import rearrange +import comfy.model_management as model_management +from comfy.ldm.modules import attention as comfy_attention +from comfy.ldm.flux import math as flux_math +from comfy.ldm.flux import layers as flux_layers +from xformers.ops import memory_efficient_attention as xattention +import numpy as np +from PIL import Image, ImageFilter, ImageDraw +from functools import partial + +class FluxAttentionControl: + def __init__(self): + self.original_attention = comfy_attention.optimized_attention + self.original_flux_attention = flux_math.attention + self.original_flux_layers_attention = flux_layers.attention + print("FluxAttentionControl initialized") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "condition": ("CONDITIONING",), + "latent_dimensions": ("LATENT",), + "region1": ("REGION",), + "number_of_regions": ("INT", { + "default": 1, + "min": 1, + "max": 3, + "step": 1, + "display": "Number of Regions" + }), + "enabled": ("BOOLEAN", { + "default": True, + "display": "Enable Regional Control" + }), + "feather_radius1": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 100.0, + "step": 1.0, + "display": "Feather Radius for Region 1" + }), + }, + "optional": { + "region2": ("REGION",), + "feather_radius2": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 100.0, + "step": 1.0, + "display": "Feather Radius for Region 2" + }), + "region3": ("REGION",), + "feather_radius3": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 100.0, + "step": 1.0, + "display": "Feather Radius for Region 3" + }), + } + } + + RETURN_TYPES = ("MODEL", "CONDITIONING",) + RETURN_NAMES = ("model", "conditioning",) + FUNCTION = "apply_attention_control" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def generate_region_mask(self, region: Dict, width: int, height: int, feather_radius: float) -> Image.Image: + if region.get('bbox') is not None: + x1, y1, x2, y2 = region['bbox'] + x1_px = int(x1 * width) + y1_px = int(y1 * height) + x2_px = int(x2 * width) + y2_px = int(y2 * height) + mask = Image.new('L', (width, height), 0) + mask_draw = ImageDraw.Draw(mask) + mask_draw.rectangle([x1_px, y1_px, x2_px, y2_px], fill=255) + if feather_radius > 0: + mask = mask.filter(ImageFilter.GaussianBlur(radius=feather_radius)) + print(f'Generating masks with {width}x{height} and [{x1}, {y1}, {x2}, {y2}], feather_radius={feather_radius}') + return mask + elif region.get('mask') is not None: + mask = region['mask'][0].cpu().numpy() + mask = (mask * 255).astype(np.uint8) + mask = Image.fromarray(mask) + mask = mask.resize((width, height)) + if feather_radius > 0: + mask = mask.filter(ImageFilter.GaussianBlur(radius=feather_radius)) + return mask + else: + raise Exception('Unknown region type') + + def generate_test_mask(self, masks: List[Image.Image], height: int, width: int): + hH, hW = int(height) // 16, int(width) // 16 + print(f'{width} {height} -> {hW} {hH}') + + lin_masks = [] + for mask in masks: + mask = mask.convert('L') + mask = torch.tensor(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255.0 # Normalize to 0-1 + mask = F.interpolate(mask, (hH, hW), mode='bilinear', align_corners=False).flatten() + lin_masks.append(mask) + return lin_masks, hH, hW + + def prepare_attention_mask(self, lin_masks: List[torch.Tensor], region_strengths: List[float], Nx: int, emb_size: int, emb_len: int): + """Prepare attention mask for three regions with per-region strengths.""" + total_len = emb_len + Nx + n_regs = len(lin_masks) + + # Initialize attention mask and scales + cross_mask = torch.zeros(total_len, total_len) + q_scale = torch.ones(total_len) + k_scale = torch.ones(total_len) + + # Indices for embeddings + main_prompt_start = 0 + main_prompt_end = emb_size + + # Subprompt indices + subprompt_starts = [emb_size * (i + 1) for i in range(n_regs)] + subprompt_ends = [emb_size * (i + 2) for i in range(n_regs)] + + # Initialize position masks + position_masks = torch.stack(lin_masks) # Shape: [n_regs, Nx] + + # Normalize masks so that overlapping areas sum to 1 + position_masks_sum = position_masks.sum(dim=0) + position_masks_normalized = position_masks / (position_masks_sum + 1e-8) + + # Build attention masks and scales + for i in range(n_regs): + sp_start = subprompt_starts[i] + sp_end = subprompt_ends[i] + mask_i = position_masks_normalized[i] + + # Scale embeddings based on mask and per-region strength + strength = region_strengths[i] + q_scale[sp_start:sp_end] = mask_i.mean() * strength + k_scale[sp_start:sp_end] = mask_i.mean() * strength + + # Create mask including tokens and positions + m_with_tokens = torch.cat([torch.ones(emb_len), mask_i]) + mb = m_with_tokens > 0.0 # Include positions where mask > 0 + + # Block attention between positions not in mask and subprompt + cross_mask[~mb, sp_start:sp_end] = 1 + cross_mask[sp_start:sp_end, ~mb] = 1 + + # Block attention between positions in region and main prompt + positions_idx = (mask_i > 0.0).nonzero(as_tuple=True)[0] + emb_len + cross_mask[positions_idx[:, None], main_prompt_start:main_prompt_end] = 1 + cross_mask[main_prompt_start:main_prompt_end, positions_idx[None, :]] = 1 + + # Block attention between subprompts + for j in range(n_regs): + if i != j: + other_sp_start = subprompt_starts[j] + other_sp_end = subprompt_ends[j] + cross_mask[sp_start:sp_end, other_sp_start:other_sp_end] = 1 + cross_mask[other_sp_start:other_sp_end, sp_start:sp_end] = 1 + + # Ensure self-attention is allowed + cross_mask.fill_diagonal_(0) + + # Prepare scales for GPU + q_scale = q_scale.reshape(1, 1, -1, 1).cuda() + k_scale = k_scale.reshape(1, 1, -1, 1).cuda() + + return cross_mask, q_scale, k_scale + + def xformers_attention(self, q: Tensor, k: Tensor, v: Tensor, pe: Tensor, + attn_mask: Optional[Tensor] = None) -> Tensor: + q, k = flux_math.apply_rope(q, k, pe) + q = rearrange(q, "B H L D -> B L H D") + k = rearrange(k, "B H L D -> B L H D") + v = rearrange(v, "B H L D -> B L H D") + if attn_mask is not None: + x = xattention(q, k, v, attn_bias=attn_mask) + else: + x = xattention(q, k, v) + x = rearrange(x, "B L H D -> B L (H D)") + return x + + def apply_attention_control(self, + model: object, + condition: List, + latent_dimensions: Dict, + region1: Dict, + number_of_regions: int, + enabled: bool, + feather_radius1: float = 0.0, + region2: Optional[Dict] = None, + feather_radius2: Optional[float] = 0.0, + region3: Optional[Dict] = None, + feather_radius3: Optional[float] = 0.0): + + # Extract dimensions and embeddings first (moved before enabled check) + latent = latent_dimensions["samples"] + bs_l, n_ch, lH, lW = latent.shape + text_emb = condition[0][0].clone() + clip_emb = condition[0][1]['pooled_output'].clone() + bs, emb_size, emb_dim = text_emb.shape + iH, iW = lH * 8, lW * 8 + + if not enabled: + # Restore original attention functions + flux_math.attention = self.original_flux_attention + flux_layers.attention = self.original_flux_layers_attention + print("Regional control disabled. Restored original attention functions.") + return (model, condition) # Return original condition when disabled + + print(f'Region attention Node enabled: {enabled}, regions: {number_of_regions}') + + # Extract dimensions and embeddings + latent = latent_dimensions["samples"] + bs_l, n_ch, lH, lW = latent.shape + text_emb = condition[0][0].clone() + clip_emb = condition[0][1]['pooled_output'].clone() + bs, emb_size, emb_dim = text_emb.shape + iH, iW = lH * 8, lW * 8 + + # Process active regions + subprompts_embeds = [] + masks = [] + region_strengths = [] + + # Collect regions and feather radii + regions = [region1, region2, region3] + feather_radii = [feather_radius1, feather_radius2, feather_radius3] + + for idx, region in enumerate(regions[:number_of_regions]): + if region is not None and region.get('conditioning') is not None: + # Get 'strength' from region or default to 1.0 + strength = region.get('strength', 1.0) + region_strengths.append(strength) + subprompt_emb = region['conditioning'][0][0] + subprompts_embeds.append(subprompt_emb) + # Use per-region feather_radius + feather_radius = feather_radii[idx] if feather_radii[idx] is not None else 0.0 + masks.append(self.generate_region_mask(region, iW, iH, feather_radius)) + else: + print(f"Region {idx+1} is None or has no conditioning") + + if not subprompts_embeds: + print("No active regions with conditioning found.") + # Restore original attention functions + flux_math.attention = self.original_flux_attention + flux_layers.attention = self.original_flux_layers_attention + return (model, condition) + + n_regs = len(subprompts_embeds) + + # Generate attention components + lin_masks, hH, hW = self.generate_test_mask(masks, iH, iW) + Nx = int(hH * hW) + emb_len = emb_size * (n_regs + 1) # +1 for main prompt + + # Create attention mask + attn_mask, q_scale, k_scale = self.prepare_attention_mask( + lin_masks, region_strengths, Nx, emb_size, emb_len) + + # Format for xFormers + device = torch.device('cuda') + attn_dtype = torch.bfloat16 if model_management.should_use_bf16(device=device) else torch.float16 + + if attn_mask is not None: + print(f'Applying attention masks: torch.Size([{attn_mask.shape[0]}, {attn_mask.shape[1]}])') + L = attn_mask.shape[0] + H = 24 # Number of heads in FLUX model + pad = (8 - L % 8) % 8 # Ensure pad is between 0 and 7 + pad_L = L + pad + mask_out = torch.zeros([bs, H, pad_L, pad_L], dtype=attn_dtype, device=device) + mask_out[:, :, :L, :L] = attn_mask.to(device, dtype=attn_dtype) + attn_mask = mask_out[:, :, :pad_L, :pad_L] + + # Prepare final mask + attn_mask_bool = attn_mask > 0.5 + attn_mask.masked_fill_(attn_mask_bool, float('-inf')) + + # Override attention + attn_mask_arg = attn_mask if enabled else None + override_attention = partial(self.xformers_attention, attn_mask=attn_mask_arg) + flux_math.attention = override_attention + flux_layers.attention = override_attention + + # Create extended conditioning + extended_condition = torch.cat([text_emb] + subprompts_embeds, dim=1) + + return (model, [[extended_condition, {'pooled_output': clip_emb}]]) + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "FluxAttentionControl": FluxAttentionControl +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxAttentionControl": "Flux Attention Control" +} \ No newline at end of file diff --git a/flux_resolution_cal_node.py b/flux_resolution_cal_node.py index 1c03ba8..ccd17b6 100644 --- a/flux_resolution_cal_node.py +++ b/flux_resolution_cal_node.py @@ -1,3 +1,11 @@ +from PIL import Image, ImageDraw, ImageFont +import numpy as np +import torch + +def pil2tensor(image): + """Convert PIL image to tensor in the correct format""" + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + class FluxResolutionNode: @classmethod def INPUT_TYPES(cls): @@ -18,19 +26,90 @@ class FluxResolutionNode: } } - RETURN_TYPES = ("INT", "INT", "STRING") - RETURN_NAMES = ("width", "height", "resolution") + RETURN_TYPES = ("INT", "INT", "STRING", "IMAGE") + RETURN_NAMES = ("width", "height", "resolution", "preview") FUNCTION = "calculate_dimensions" CATEGORY = "ControlAltAI Nodes/Flux" OUTPUT_NODE = True + def create_preview_image(self, width, height, resolution, ratio_display): + # 1024x1024 preview size + preview_size = (1024, 1024) + image = Image.new('RGB', preview_size, (0, 0, 0)) # Black background + draw = ImageDraw.Draw(image) + + # Draw grid with grey lines + grid_color = '#333333' # Dark grey for grid + grid_spacing = 50 # Adjusted grid spacing + for x in range(0, preview_size[0], grid_spacing): + draw.line([(x, 0), (x, preview_size[1])], fill=grid_color) + for y in range(0, preview_size[1], grid_spacing): + draw.line([(0, y), (preview_size[0], y)], fill=grid_color) + + # Calculate preview box dimensions + preview_width = 800 # Increased size + preview_height = int(preview_width * (height / width)) + + # Adjust if height is too tall + if preview_height > 800: # Adjusted for larger preview + preview_height = 800 + preview_width = int(preview_height * (width / height)) + + # Calculate center position + x_offset = (preview_size[0] - preview_width) // 2 + y_offset = (preview_size[1] - preview_height) // 2 + + # Draw the aspect ratio box with thicker outline + draw.rectangle( + [(x_offset, y_offset), (x_offset + preview_width, y_offset + preview_height)], + outline='red', + width=4 # Thicker outline + ) + + # Add text with larger font sizes + try: + # Draw text (centered) + text_y = y_offset + preview_height//2 + + # Resolution text in red + draw.text((preview_size[0]//2, text_y), + f"{width}x{height}", + fill='red', + anchor="mm", + font=ImageFont.truetype("arial.ttf", 48)) + + # Aspect ratio text in red + draw.text((preview_size[0]//2, text_y + 60), + f"({ratio_display})", + fill='red', + anchor="mm", + font=ImageFont.truetype("arial.ttf", 36)) + + # Resolution text at bottom in white + draw.text((preview_size[0]//2, y_offset + preview_height + 60), + f"Resolution: {resolution}", + fill='white', # Changed to white + anchor="mm", + font=ImageFont.truetype("arial.ttf", 32)) + + except: + # Fallback if font loading fails + draw.text((preview_size[0]//2, text_y), f"{width}x{height}", fill='red', anchor="mm") + draw.text((preview_size[0]//2, text_y + 60), f"({ratio_display})", fill='red', anchor="mm") + draw.text((preview_size[0]//2, y_offset + preview_height + 60), f"Resolution: {resolution}", fill='white', anchor="mm") + + # Convert to tensor using the helper function + return pil2tensor(image) + def calculate_dimensions(self, megapixel, aspect_ratio, custom_ratio, custom_aspect_ratio=None): megapixel = float(megapixel) if custom_ratio and custom_aspect_ratio: numeric_ratio = custom_aspect_ratio + ratio_display = custom_aspect_ratio # Keep original format for display else: numeric_ratio = aspect_ratio.split(' ')[0] + ratio_display = numeric_ratio # Keep original format for display width_ratio, height_ratio = map(int, numeric_ratio.split(':')) @@ -52,7 +131,10 @@ class FluxResolutionNode: resolution = f"{width} x {height}" - return width, height, resolution + # Generate preview image with original ratio format + preview = self.create_preview_image(width, height, resolution, ratio_display) + + return width, height, resolution, preview NODE_CLASS_MAPPINGS = { "FluxResolutionNode": FluxResolutionNode, diff --git a/region_mask_conditioning_node.py b/region_mask_conditioning_node.py new file mode 100644 index 0000000..0c77b80 --- /dev/null +++ b/region_mask_conditioning_node.py @@ -0,0 +1,301 @@ +import torch +import numpy as np +from typing import Dict, List, Optional, Tuple +from PIL import Image, ImageDraw + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +class RegionMaskConditioning: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask1": ("MASK",), + "bbox1": ("BBOX",), + "conditioning1": ("CONDITIONING",), + "number_of_regions": ("INT", { + "default": 1, + "min": 1, + "max": 3, + "step": 1, + "display": "Number of Regions" + }), + "strength1": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.1, + "display": "Strength for Region 1" + }), + }, + "optional": { + "mask2": ("MASK",), + "bbox2": ("BBOX",), + "conditioning2": ("CONDITIONING",), + "strength2": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.1, + "display": "Strength for Region 2" + }), + "mask3": ("MASK",), + "bbox3": ("BBOX",), + "conditioning3": ("CONDITIONING",), + "strength3": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.1, + "display": "Strength for Region 3" + }), + } + } + + RETURN_TYPES = ("REGION", "REGION", "REGION", "INT", "IMAGE") + RETURN_NAMES = ("region1", "region2", "region3", + "region_count", "preview_image") + FUNCTION = "create_conditioned_regions" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def validate_bbox(self, bbox: Dict) -> bool: + """Validate bbox coordinates and structure""" + print(f"\n=== Validating BBox ===") + print(f"Input bbox: {bbox}") + + if bbox is None or not isinstance(bbox, dict): + print("Failed: Invalid bbox type") + return False + + required_keys = ["x1", "y1", "x2", "y2"] + if not all(k in bbox for k in required_keys): + print(f"Failed: Missing keys. Required: {required_keys}") + return False + + # Validate coordinate values + if not all(isinstance(bbox[k], (int, float)) for k in required_keys): + print("Failed: Invalid coordinate types") + return False + + # Validate coordinate ranges + if not all(0 <= bbox[k] <= 1.0 for k in required_keys): + print("Failed: Coordinates out of range [0,1]") + return False + + # Validate proper ordering + if bbox["x1"] >= bbox["x2"] or bbox["y1"] >= bbox["y2"]: + print("Failed: Invalid coordinate ordering") + return False + + print("Passed: BBox validation successful") + return True + + def scale_conditioning(self, conditioning: List, strength: float) -> List: + """Scale conditioning tensors by strength""" + print(f"\n=== Scaling Conditioning ===") + print(f"Strength: {strength}") + + try: + if not conditioning or not isinstance(conditioning, list): + print("Failed: Invalid conditioning format") + raise ValueError("Invalid conditioning format") + + # Get the conditioning tensors and dict + cond_tensors = conditioning[0][0] + cond_dict = conditioning[0][1] + + print(f"Input tensor shape: {cond_tensors.shape}") + print(f"Conditioning keys: {list(cond_dict.keys())}") + print(f"Input tensor stats: min={cond_tensors.min():.3f}, max={cond_tensors.max():.3f}, mean={cond_tensors.mean():.3f}") + + # Scale the tensors + scaled_tensors = cond_tensors.clone() * strength + print(f"Scaled tensor stats: min={scaled_tensors.min():.3f}, max={scaled_tensors.max():.3f}, mean={scaled_tensors.mean():.3f}") + + return [[scaled_tensors, cond_dict]] + + except Exception as e: + print(f"Error in scale_conditioning: {str(e)}") + import traceback + traceback.print_exc() + return conditioning + + def create_region(self, mask: Optional[torch.Tensor], bbox: Optional[Dict], + conditioning: Optional[List], strength: float, region_idx: int) -> Dict: + """Create a single region with its conditioning""" + print(f"\n=== Creating Region {region_idx} ===") + + # Debug inputs + print("Input validation:") + print(f"- Mask: {type(mask)}, shape={mask.shape if mask is not None else None}") + print(f"- BBox: {bbox}") + print(f"- Conditioning type: {type(conditioning)}") + print(f"- Strength: {strength}") + + # Default empty region + empty_region = { + "conditioning": None, + "bbox": [0.0, 0.0, 0.0, 0.0], # Array format for empty + "is_active": False, + "strength": 1.0 + } + + try: + # Validate inputs + if mask is None or bbox is None or conditioning is None: + print(f"Region {region_idx}: Missing components") + return empty_region + + if not self.validate_bbox(bbox): + print(f"Region {region_idx}: Invalid bbox") + return empty_region + + # Scale conditioning + scaled_conditioning = self.scale_conditioning(conditioning, strength) + + # Create region output - bbox array, conditioning, and strength + region = { + "conditioning": scaled_conditioning, + "bbox": [bbox["x1"], bbox["y1"], bbox["x2"], bbox["y2"]], # Array format + "is_active": True, + "strength": strength + } + + print(f"\nSuccessfully created region {region_idx}") + return region + + except Exception as e: + print(f"Error creating region {region_idx}: {str(e)}") + import traceback + traceback.print_exc() + return empty_region + + def create_preview(self, masks: List[torch.Tensor], bboxes: List[Dict], + number_of_regions: int) -> torch.Tensor: + """Create preview of conditioned regions""" + print("\n=== Creating Preview ===") + + if not masks: + print("No masks provided") + return torch.zeros((3, 64, 64), dtype=torch.float32) + + height, width = masks[0].shape + print(f"Preview dimensions: {width}x{height}") + + # Create PIL Image for preview + preview = Image.new("RGB", (width, height), (0, 0, 0)) + draw = ImageDraw.Draw(preview) + + # Define colors for 3 regions + colors = [ + (255, 0, 0), # Red + (0, 255, 0), # Green + (255, 255, 0), # Yellow + ] + + # Draw each region + for i, (mask, bbox) in enumerate(zip(masks[:number_of_regions], bboxes[:number_of_regions])): + validation_result = self.validate_bbox(bbox) + if validation_result and mask is not None: + print(f"\nDrawing region {i+1}:") + # Get pixel coordinates + x1 = int(bbox["x1"] * width) + y1 = int(bbox["y1"] * height) + x2 = int(bbox["x2"] * width) + y2 = int(bbox["y2"] * height) + + print(f"Region {i+1} coordinates: ({x1},{y1}) to ({x2},{y2})") + + # Draw region outline + draw.rectangle([x1, y1, x2, y2], outline=colors[i], width=4) + + return pil2tensor(preview) + + def create_conditioned_regions(self, + mask1: torch.Tensor, + bbox1: Dict, + conditioning1: List, + number_of_regions: int, + strength1: float, + mask2: Optional[torch.Tensor] = None, + bbox2: Optional[Dict] = None, + conditioning2: Optional[List] = None, + strength2: Optional[float] = 1.0, + mask3: Optional[torch.Tensor] = None, + bbox3: Optional[Dict] = None, + conditioning3: Optional[List] = None, + strength3: Optional[float] = 1.0) -> Tuple: + print("\n=== Creating Conditioned Regions ===") + print(f"Number of regions: {number_of_regions}") + + try: + # Create regions + regions = [] + active_count = 0 + + # Process required number of regions + inputs = [ + (mask1, bbox1, conditioning1, strength1), + (mask2, bbox2, conditioning2, strength2), + (mask3, bbox3, conditioning3, strength3) + ] + + # Store masks and bboxes for preview only + preview_masks = [] + preview_bboxes = [] + + for i, (mask, bbox, conditioning, strength) in enumerate(inputs[:number_of_regions]): + # Create region with per-region strength + region = self.create_region(mask, bbox, conditioning, strength, i+1) + if region["is_active"]: + active_count += 1 + regions.append(region) + print(f"Processed region {i+1}: active={region['is_active']}") + + # Store for preview + preview_masks.append(mask) + preview_bboxes.append(bbox) + + # Fill remaining slots with empty regions + empty_region = { + "conditioning": None, + "bbox": [0.0, 0.0, 0.0, 0.0], # Array format + "is_active": False, + "strength": 1.0 + } + + while len(regions) < 3: + idx = len(regions) + 1 + print(f"Adding empty region {idx}") + regions.append(empty_region) + + print(f"\nCreated {active_count} active regions out of {number_of_regions} requested") + + # Create preview using stored masks and bboxes + preview = self.create_preview(preview_masks, preview_bboxes, number_of_regions) + + return (*regions, active_count, preview) + + except Exception as e: + print(f"Error in create_conditioned_regions: {str(e)}") + import traceback + traceback.print_exc() + + empty_region = { + "conditioning": None, + "bbox": [0.0, 0.0, 0.0, 0.0], # Array format + "is_active": False, + "strength": 1.0 + } + empty_preview = torch.zeros((3, mask1.shape[0], mask1.shape[1]), dtype=torch.float32) + return (empty_region, empty_region, empty_region, 0, empty_preview) + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "RegionMaskConditioning": RegionMaskConditioning +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RegionMaskConditioning": "Region Mask Conditioning" +} \ No newline at end of file diff --git a/region_mask_generator_node.py b/region_mask_generator_node.py new file mode 100644 index 0000000..a732a42 --- /dev/null +++ b/region_mask_generator_node.py @@ -0,0 +1,164 @@ +import torch +import numpy as np +from typing import Dict, List, Optional, Tuple +from PIL import Image, ImageDraw + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +class RegionMaskGenerator: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "width": ("INT", {"default": 1024}), + "height": ("INT", {"default": 1024}), + "number_of_regions": ("INT", {"default": 1, "min": 1, "max": 3}), + # Region 1 + "region1_x1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "region1_y1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "region1_x2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "region1_y2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + # Region 2 + "region2_x1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "region2_y1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "region2_x2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "region2_y2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + # Region 3 + "region3_x1": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "region3_y1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "region3_x2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "region3_y2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "MASK", "MASK", "INT", "BBOX", "BBOX", "BBOX") + RETURN_NAMES = ("colored_regions_image", "bbox_preview", + "mask1", "mask2", "mask3", + "number_of_regions", + "bbox1", "bbox2", "bbox3") + FUNCTION = "generate_regions" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def create_bbox(self, x1: float, y1: float, x2: float, y2: float) -> Dict: + """Create bbox with debug output""" + print(f"Creating BBOX: x1={x1:.3f}, y1={y1:.3f}, x2={x2:.3f}, y2={y2:.3f}") + return { + "x1": x1, + "y1": y1, + "x2": x2, + "y2": y2, + "active": True + } + + def create_mask_from_bbox(self, bbox: Dict, width: int, height: int) -> torch.Tensor: + """Create mask from bbox with debug output""" + mask = torch.zeros((height, width), dtype=torch.float32) + if bbox["active"]: + x1 = int(bbox["x1"] * width) + y1 = int(bbox["y1"] * height) + x2 = int(bbox["x2"] * width) + y2 = int(bbox["y2"] * height) + print(f"Creating mask at pixels: x1={x1}, y1={y1}, x2={x2}, y2={y2}") + mask[y1:y2, x1:x2] = 1.0 + return mask + + def create_preview(self, masks: List[torch.Tensor], bboxes: List[Dict], + number_of_regions: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Create preview images with debug info""" + if not masks: + return torch.zeros((3, 64, 64), dtype=torch.float32), torch.zeros((3, 64, 64), dtype=torch.float32) + + height, width = masks[0].shape + + # Create both preview images + region_preview = Image.new("RGB", (width, height), (0, 0, 0)) + bbox_preview = Image.new("RGB", (width, height), (0, 0, 0)) + region_draw = ImageDraw.Draw(region_preview) + bbox_draw = ImageDraw.Draw(bbox_preview) + + colors = [ + (255, 0, 0), # Red - Region 1 + (0, 255, 0), # Green - Region 2 + (255, 255, 0), # Yellow - Region 3 + ] + + # Store regions for ordered preview + preview_regions = [] + for i in range(number_of_regions): + if bboxes[i]["active"]: + mask_np = masks[i].cpu().numpy() > 0.5 + if mask_np.any(): + preview_regions.append((i, mask_np, bboxes[i])) + + # Draw regions in reverse order (Region 3 first, Region 1 last) + for i, mask_np, bbox in sorted(preview_regions, reverse=True): + # Draw on region preview + color_array = np.zeros((height, width, 3), dtype=np.uint8) + for c in range(3): + color_array[mask_np, c] = colors[i][c] + preview_np = np.array(region_preview) + preview_np[mask_np] = color_array[mask_np] + region_preview = Image.fromarray(preview_np) + + # Draw on bbox preview - maintaining original bbox drawing order + x1 = int(bbox["x1"] * width) + y1 = int(bbox["y1"] * height) + x2 = int(bbox["x2"] * width) + y2 = int(bbox["y2"] * height) + print(f"Drawing preview for region {i}: x1={x1}, y1={y1}, x2={x2}, y2={y2}") + bbox_draw.rectangle([x1, y1, x2, y2], outline=colors[i], width=2) + + return pil2tensor(region_preview), pil2tensor(bbox_preview) + + def generate_regions(self, + width: int, + height: int, + number_of_regions: int, + **kwargs) -> Tuple: + try: + print(f"\nGenerating {number_of_regions} regions for {width}x{height} image") + bboxes = [] + masks = [] + + # Create regions + for i in range(3): + if i < number_of_regions: + print(f"\nProcessing region {i+1}:") + bbox = self.create_bbox( + kwargs[f"region{i+1}_x1"], + kwargs[f"region{i+1}_y1"], + kwargs[f"region{i+1}_x2"], + kwargs[f"region{i+1}_y2"] + ) + mask = self.create_mask_from_bbox(bbox, width, height) + bboxes.append(bbox) + masks.append(mask) + else: + print(f"Creating empty region {i+1}") + empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False} + bboxes.append(empty_bbox) + masks.append(torch.zeros((height, width), dtype=torch.float32)) + + # Create previews + region_preview, bbox_preview = self.create_preview(masks, bboxes, number_of_regions) + + return (region_preview, bbox_preview, *masks, number_of_regions, *bboxes) + + except Exception as e: + print(f"Error in generate_regions: {str(e)}") + empty_mask = torch.zeros((height, width), dtype=torch.float32) + empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False} + empty_preview = torch.zeros((3, height, width), dtype=torch.float32) + return (empty_preview, empty_preview, + empty_mask, empty_mask, empty_mask, + 0, empty_bbox, empty_bbox, empty_bbox) + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "RegionMaskGenerator": RegionMaskGenerator +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RegionMaskGenerator": "Region Mask Generator" +} \ No newline at end of file diff --git a/region_mask_processor_node.py b/region_mask_processor_node.py new file mode 100644 index 0000000..e2f0e54 --- /dev/null +++ b/region_mask_processor_node.py @@ -0,0 +1,230 @@ +import torch +import torch.nn.functional as F +from typing import Tuple, Dict, Optional, List +import numpy as np +from PIL import Image, ImageDraw + +def pil2tensor(image): + """Convert a PIL image to a PyTorch tensor in the expected format.""" + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +class RegionMaskProcessor: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask1": ("MASK",), + "bbox1": ("BBOX",), + "blur_radius": ("INT", { + "default": 5, + "min": 0, + "max": 32, + "step": 1, + "display": "Blur Radius" + }), + "threshold": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.1, + "display": "Mask Threshold" + }), + "feather_edges": ("BOOLEAN", { + "default": True, + "display": "Feather Edges" + }), + "number_of_regions": ("INT", { + "default": 1, + "min": 1, + "max": 3, + "display": "Number of Regions" + }), + }, + "optional": { + "mask2": ("MASK",), + "bbox2": ("BBOX",), + "mask3": ("MASK",), + "bbox3": ("BBOX",), + } + } + + RETURN_TYPES = ("MASK", "BBOX", "MASK", "BBOX", "MASK", "BBOX", "IMAGE", "INT") + RETURN_NAMES = ("processed_mask1", "processed_bbox1", + "processed_mask2", "processed_bbox2", + "processed_mask3", "processed_bbox3", + "preview_image", "region_count") + FUNCTION = "process_regions" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def apply_gaussian_blur(self, mask: torch.Tensor, radius: int) -> torch.Tensor: + """Apply gaussian blur to mask edges""" + if radius <= 0: + return mask + + kernel_size = 2 * radius + 1 + sigma = radius / 3.0 + + if len(mask.shape) == 2: + mask = mask.unsqueeze(0).unsqueeze(0) + + kernel_1d = torch.exp(torch.linspace(-radius, radius, kernel_size).pow(2) / (-2 * sigma ** 2)) + kernel_1d = kernel_1d / kernel_1d.sum() + + padding = radius + kernel_h = kernel_1d.unsqueeze(0).unsqueeze(0).unsqueeze(0).to(mask.device) + kernel_v = kernel_1d.unsqueeze(0).unsqueeze(0).unsqueeze(-1).to(mask.device) + + mask = F.pad(mask, (padding, padding, 0, 0), mode='reflect') + mask = F.conv2d(mask, kernel_h) + mask = F.pad(mask, (0, 0, padding, padding), mode='reflect') + mask = F.conv2d(mask, kernel_v) + + return mask.squeeze() + + def apply_feathering(self, mask: torch.Tensor, bbox: Dict, radius: int) -> Tuple[torch.Tensor, Dict]: + """Apply feathering to mask edges while preserving bbox boundaries""" + if radius <= 0 or not bbox["active"]: + return mask, bbox + + height, width = mask.shape + x1 = int(bbox["x1"] * width) + y1 = int(bbox["y1"] * height) + x2 = int(bbox["x2"] * width) + y2 = int(bbox["y2"] * height) + + inner_mask = torch.zeros_like(mask) + inner_mask[y1+radius:y2-radius, x1+radius:x2-radius] = 1.0 + edge_mask = mask - inner_mask + + if edge_mask.any(): + blurred = self.apply_gaussian_blur(mask, radius) + result = mask.clone() + result[edge_mask > 0] = blurred[edge_mask > 0] + else: + result = mask + + return result, bbox + + def process_single_region(self, + mask: torch.Tensor, + bbox: Dict, + blur_radius: int, + threshold: float, + feather_edges: bool) -> Tuple[torch.Tensor, Dict]: + """Process a single mask-bbox pair""" + if mask is None or not bbox["active"]: + return mask, bbox + + try: + processed = (mask > threshold).float() + + if feather_edges and blur_radius > 0: + processed, bbox = self.apply_feathering(processed, bbox, blur_radius) + elif blur_radius > 0: + processed = self.apply_gaussian_blur(processed, blur_radius) + + return processed, bbox + + except Exception as e: + print(f"Error processing region: {str(e)}") + return mask, bbox + + def create_preview(self, masks: List[torch.Tensor], bboxes: List[Dict], + number_of_regions: int) -> torch.Tensor: + """Create preview of processed regions with PIL for consistent coloring""" + if not masks: + return torch.zeros((3, 64, 64), dtype=torch.float32) + + height, width = masks[0].shape + + # Create PIL Image for preview + preview = Image.new("RGB", (width, height), (0, 0, 0)) + + colors = [ + (255, 0, 0), # Red - Region 1 + (0, 255, 0), # Green - Region 2 + (255, 255, 0), # Yellow - Region 3 + ] + + # Store regions for ordered preview + preview_regions = [] + for i in range(number_of_regions): + if bboxes[i]["active"] and masks[i] is not None: + mask_np = masks[i].cpu().numpy() > 0.5 + preview_regions.append((i, mask_np)) + + # Draw regions in reverse order (Region 3 first, Region 1 last) + for i, mask_np in sorted(preview_regions, reverse=True): + color_array = np.zeros((height, width, 3), dtype=np.uint8) + color_array[mask_np] = colors[i] + + # Convert to PIL and composite + region_img = Image.fromarray(color_array, 'RGB') + preview = Image.alpha_composite( + preview.convert('RGBA'), + Image.merge('RGBA', (*region_img.split(), Image.fromarray((mask_np * 255).astype(np.uint8)))) + ) + + return pil2tensor(preview.convert('RGB')) + + def process_regions(self, + mask1: torch.Tensor, + bbox1: Dict, + blur_radius: int, + threshold: float, + feather_edges: bool, + number_of_regions: int, + mask2: Optional[torch.Tensor] = None, + bbox2: Optional[Dict] = None, + mask3: Optional[torch.Tensor] = None, + bbox3: Optional[Dict] = None) -> Tuple: + try: + # Process each mask-bbox pair + mask_bbox_pairs = [ + (mask1, bbox1), + (mask2, bbox2) if mask2 is not None else (None, None), + (mask3, bbox3) if mask3 is not None else (None, None), + ] + + processed_masks = [] + processed_bboxes = [] + active_count = 0 + + for i, (mask, bbox) in enumerate(mask_bbox_pairs): + if i < number_of_regions and mask is not None and bbox is not None: + proc_mask, proc_bbox = self.process_single_region( + mask, bbox, blur_radius, threshold, feather_edges + ) + if proc_bbox["active"]: + active_count += 1 + processed_masks.append(proc_mask) + processed_bboxes.append(proc_bbox) + else: + empty_mask = torch.zeros_like(mask1) + empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False} + processed_masks.append(empty_mask) + processed_bboxes.append(empty_bbox) + + # Create preview + preview = self.create_preview(processed_masks, processed_bboxes, number_of_regions) + + return (*[item for pair in zip(processed_masks, processed_bboxes) for item in pair], + preview, active_count) + + except Exception as e: + print(f"Error processing regions: {str(e)}") + empty_mask = torch.zeros_like(mask1) + empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False} + empty_preview = torch.zeros((3, mask1.shape[0], mask1.shape[1]), dtype=torch.float32) + return (empty_mask, empty_bbox, empty_mask, empty_bbox, + empty_mask, empty_bbox, + empty_preview, 0) + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "RegionMaskProcessor": RegionMaskProcessor +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RegionMaskProcessor": "Region Mask Processor" +} \ No newline at end of file diff --git a/region_mask_validator_node.py b/region_mask_validator_node.py new file mode 100644 index 0000000..b08d96f --- /dev/null +++ b/region_mask_validator_node.py @@ -0,0 +1,270 @@ +import torch +from typing import Tuple, Dict, Optional, List +import numpy as np +from PIL import Image, ImageDraw + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +class RegionMaskValidator: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask1": ("MASK",), + "bbox1": ("BBOX",), + "number_of_regions": ("INT", { + "default": 1, + "min": 1, + "max": 3, + "step": 1 + }), + "min_region_size": ("INT", { + "default": 64, + "min": 32, + "max": 512, + "step": 32, + "display": "Minimum Region Size (px)" + }), + "max_overlap": ("FLOAT", { + "default": 0.1, + "min": 0.0, + "max": 0.5, + "step": 0.01, + "display": "Maximum Region Overlap" + }), + }, + "optional": { + "mask2": ("MASK",), + "bbox2": ("BBOX",), + "mask3": ("MASK",), + "bbox3": ("BBOX",), + } + } + + RETURN_TYPES = ("MASK", "BBOX", "MASK", "BBOX", "MASK", "BBOX", + "INT", "BOOLEAN", "STRING", "IMAGE") + RETURN_NAMES = ("valid_mask1", "valid_bbox1", + "valid_mask2", "valid_bbox2", + "valid_mask3", "valid_bbox3", + "valid_region_count", "is_valid", "validation_message", + "validation_preview") + FUNCTION = "validate_regions" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def get_region_dimensions(self, bbox: Dict, width: int, height: int) -> Tuple[int, int, Tuple[int, int]]: + """Calculate region dimensions in pixels""" + if not bbox["active"]: + return 0, (0, 0) + + x1 = int(bbox["x1"] * width) + y1 = int(bbox["y1"] * height) + x2 = int(bbox["x2"] * width) + y2 = int(bbox["y2"] * height) + + w = x2 - x1 + h = y2 - y1 + area = w * h + print(f"Region dimensions: {w}x{h} pixels") + return area, (w, h) + + def calculate_overlap(self, bbox1: Dict, bbox2: Dict, width: int, height: int) -> Tuple[Tuple[int, int], float]: + """Calculate overlap dimensions and ratio""" + if not (bbox1["active"] and bbox2["active"]): + return (0, 0), 0.0 + + # Convert to pixel coordinates + x1_1 = int(bbox1["x1"] * width) + y1_1 = int(bbox1["y1"] * height) + x2_1 = int(bbox1["x2"] * width) + y2_1 = int(bbox1["y2"] * height) + + x1_2 = int(bbox2["x1"] * width) + y1_2 = int(bbox2["y1"] * height) + x2_2 = int(bbox2["x2"] * width) + y2_2 = int(bbox2["y2"] * height) + + # Calculate intersection + x_left = max(x1_1, x1_2) + y_top = max(y1_1, y1_2) + x_right = min(x2_1, x2_2) + y_bottom = min(y2_1, y2_2) + + if x_right > x_left and y_bottom > y_top: + overlap_width = x_right - x_left + overlap_height = y_bottom - y_top + overlap_area = overlap_width * overlap_height + + area1 = (x2_1 - x1_1) * (y2_1 - y1_1) + area2 = (x2_2 - x1_2) * (y2_2 - y1_2) + smaller_area = min(area1, area2) + overlap_ratio = overlap_area / smaller_area + + print(f"Overlap dimensions: {overlap_width}x{overlap_height} pixels ({overlap_ratio:.1%})") + return (overlap_width, overlap_height), overlap_ratio + + return (0, 0), 0.0 + + def create_validation_preview(self, masks: List[torch.Tensor], bboxes: List[Dict], + number_of_regions: int, is_valid: bool, + messages: List[str], img_width: int, img_height: int) -> torch.Tensor: + """Create visual validation feedback with improved text rendering""" + if not masks: + return torch.zeros((3, 64, 64), dtype=torch.float32) + + preview = Image.new("RGB", (img_width, img_height), (0, 0, 0)) + draw = ImageDraw.Draw(preview) + + # Colors for valid/invalid regions + colors = { + 'valid': [(0, 255, 0), (0, 200, 0), (0, 150, 0)], # Green shades + 'invalid': [(255, 0, 0), (200, 0, 0), (150, 0, 0)] # Red shades + } + + # Draw regions with validation status and improved text + for i, (mask, bbox) in enumerate(zip(masks[:number_of_regions], bboxes[:number_of_regions])): + if bbox["active"]: + x1 = int(bbox["x1"] * img_width) + y1 = int(bbox["y1"] * img_height) + x2 = int(bbox["x2"] * img_width) + y2 = int(bbox["y2"] * img_height) + + w = x2 - x1 + h = y2 - y1 + color = colors['valid' if is_valid else 'invalid'][i] + + # Draw thicker rectangle outline + draw.rectangle([x1, y1, x2, y2], outline=color, width=4) + + # Improved region label with dimensions + label = f"R{i+1}: {w}x{h}" + # Position text with offset from corner and draw twice for better visibility + text_x = x1 + 10 + text_y = y1 + 10 + + # Draw text shadow/outline for better contrast + shadow_offset = 2 + shadow_color = (0, 0, 0) + for dx in [-shadow_offset, shadow_offset]: + for dy in [-shadow_offset, shadow_offset]: + draw.text((text_x + dx, text_y + dy), label, fill=shadow_color, font=None, size=64) + + # Draw main text + draw.text((text_x, text_y), label, fill=color, font=None, size=64) + + # If region is invalid, add error message below the label + if not is_valid and i < len(messages): + error_y = text_y + 30 # Position error message below label + # Draw error message with shadow for contrast + for dx in [-shadow_offset, shadow_offset]: + for dy in [-shadow_offset, shadow_offset]: + draw.text((text_x + dx, error_y + dy), messages[i], fill=shadow_color, font=None, size=20) + draw.text((text_x, error_y), messages[i], fill=color, font=None, size=20) + + return pil2tensor(preview) + + def validate_regions(self, + mask1: torch.Tensor, + bbox1: Dict, + number_of_regions: int, + min_region_size: int, + max_overlap: float, + mask2: Optional[torch.Tensor] = None, + bbox2: Optional[Dict] = None, + mask3: Optional[torch.Tensor] = None, + bbox3: Optional[Dict] = None) -> Tuple: + try: + print(f"\nValidating {number_of_regions} regions:") + messages = [] + is_valid = True + height, width = mask1.shape + print(f"Canvas size: {width}x{height} pixels") + + # Collect regions + regions = [ + (mask1, bbox1), + (mask2, bbox2) if mask2 is not None else (None, None), + (mask3, bbox3) if mask3 is not None else (None, None), + ] + + # Validate each region + valid_regions = [] + valid_count = 0 + for i, (mask, bbox) in enumerate(regions): + if i < number_of_regions and mask is not None and bbox is not None: + print(f"\nValidating Region {i+1}:") + # Check region size + _, (w, h) = self.get_region_dimensions(bbox, width, height) + + if w < min_region_size or h < min_region_size: + message = f"Region {i+1} too small: {w}x{h} pixels (minimum: {min_region_size}x{min_region_size})" + print(f"Failed: {message}") + messages.append(message) + is_valid = False + bbox = bbox.copy() + bbox["active"] = False + else: + print(f"Passed: Region {i+1} size check ({w}x{h} pixels)") + valid_count += 1 + + valid_regions.append((mask, bbox)) + else: + valid_regions.append(( + torch.zeros_like(mask1), + {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False} + )) + + # Check overlaps + if valid_count > 1: + print("\nChecking region overlaps:") + for i in range(len(valid_regions)): + for j in range(i + 1, len(valid_regions)): + mask_i, bbox_i = valid_regions[i] + mask_j, bbox_j = valid_regions[j] + + if bbox_i["active"] and bbox_j["active"]: + print(f"Checking overlap between regions {i+1} and {j+1}:") + (ow, oh), overlap_ratio = self.calculate_overlap(bbox_i, bbox_j, width, height) + + if overlap_ratio > max_overlap: + message = f"Excessive overlap ({ow}x{oh} pixels, {overlap_ratio:.1%}) between regions {i+1} and {j+1}" + print(f"Failed: {message}") + messages.append(message) + is_valid = False + + # Create validation message + validation_message = "All regions valid" if is_valid else "\n".join(messages) + print(f"\nValidation {'passed' if is_valid else 'failed'}:") + print(validation_message) + + # Create validation preview + preview = self.create_validation_preview( + [r[0] for r in valid_regions], + [r[1] for r in valid_regions], + number_of_regions, + is_valid, + messages, + width, + height + ) + + return (*[item for region in valid_regions for item in region], + valid_count, is_valid, validation_message, preview) + + except Exception as e: + print(f"Validation error: {str(e)}") + empty_mask = torch.zeros_like(mask1) + empty_bbox = {"x1": 0.0, "y1": 0.0, "x2": 0.0, "y2": 0.0, "active": False} + empty_preview = torch.zeros((3, height, width), dtype=torch.float32) + return (empty_mask, empty_bbox, empty_mask, empty_bbox, + empty_mask, empty_bbox, + 0, False, f"Validation error: {str(e)}", empty_preview) + +# Node class mappings +NODE_CLASS_MAPPINGS = { + "RegionMaskValidator": RegionMaskValidator +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RegionMaskValidator": "Region Mask Validator" +} \ No newline at end of file diff --git a/region_overlay_visualizer_node.py b/region_overlay_visualizer_node.py new file mode 100644 index 0000000..758f03f --- /dev/null +++ b/region_overlay_visualizer_node.py @@ -0,0 +1,90 @@ +import torch +import numpy as np +from typing import Tuple + +class RegionOverlayVisualizer: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "region_preview": ("IMAGE",), + "opacity": ("FLOAT", { + "default": 0.3, + "min": 0.0, + "max": 1.0, + "step": 0.1, + "display": "Overlay Opacity" + }), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "visualize_regions" + CATEGORY = "ControlAltAI Nodes/Flux Region" + + def visualize_regions( + self, + image: torch.Tensor, + region_preview: torch.Tensor, + opacity: float, + ) -> Tuple[torch.Tensor]: + try: + print("\n=== Starting Region Overlay Visualization ===") + print(f"Initial shapes - Image: {image.shape}, Preview: {region_preview.shape}") + + # Ensure input tensors are in [B, H, W, C] format + if len(image.shape) == 3: + image = image.unsqueeze(0) + if len(region_preview.shape) == 3: + region_preview = region_preview.unsqueeze(0) + + # Get working copies + base_image = image.clone() + preview = region_preview.clone() + + # Convert to numpy for mask creation (keeping batch and HWC format) + preview_np = (preview * 255).byte().cpu().numpy() + + # Create mask based on preview content (operating on the last dimension - channels) + color_sum = np.sum(preview_np, axis=-1) # Sum across color channels + max_channel = np.max(preview_np, axis=-1) + min_channel = np.min(preview_np, axis=-1) + + # Create binary mask where content exists + mask = ( + (color_sum > 50) & + (max_channel > 30) & + ((max_channel - min_channel) > 10) + ) + + # Expand mask to match input dimensions + mask = mask[..., None] # Add channel dimension back + mask = torch.from_numpy(mask).to(image.device) + + print(f"Mask shape: {mask.shape}") + print(f"Masked pixels: {mask.sum().item()}/{mask.numel()} ({mask.sum().item()/mask.numel()*100:.2f}%)") + + # Apply blending only where mask is True + result = torch.where( + mask.bool(), + (1 - opacity) * base_image + opacity * preview, + base_image + ) + + print(f"Final shape: {result.shape}") + return (result,) + + except Exception as e: + print(f"Error in visualization: {str(e)}") + import traceback + traceback.print_exc() + return (image,) + +NODE_CLASS_MAPPINGS = { + "RegionOverlayVisualizer": RegionOverlayVisualizer +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RegionOverlayVisualizer": "Region Overlay Visualizer" +} \ No newline at end of file