edit_space.py

import os
import torch.nn.functional as F
import torch
import sys
import cv2
import numpy as np
from PIL import Image
import json


# New imports for the diffuser pipeline
from src.pipeline_flux_kontext_control import FluxKontextControlPipeline
from src.transformer_flux import FluxTransformer2DModel

import tempfile
from safetensors.torch import load_file, save_file

_original_load_lora_weights = FluxKontextControlPipeline.load_lora_weights

def _patched_load_lora_weights(self, pretrained_model_name_or_path_or_dict, **kwargs):
    """自动转换混合格式的 LoRA 并添加 transformer 前缀"""
    weight_name = kwargs.get("weight_name", "pytorch_lora_weights.safetensors")
    
    if isinstance(pretrained_model_name_or_path_or_dict, str):
        if os.path.isdir(pretrained_model_name_or_path_or_dict):
            lora_file = os.path.join(pretrained_model_name_or_path_or_dict, weight_name)
        else:
            lora_file = pretrained_model_name_or_path_or_dict
        
        if os.path.exists(lora_file):
            state_dict = load_file(lora_file)
            
            # 检查是否需要转换格式或添加前缀
            needs_format_conversion = any('lora_A.weight' in k or 'lora_B.weight' in k for k in state_dict.keys())
            needs_prefix = not any(k.startswith('transformer.') for k in state_dict.keys())
            
            if needs_format_conversion or needs_prefix:
                print(f"🔄 Processing LoRA: {lora_file}")
                if needs_format_conversion:
                    print(f"   - Converting PEFT format to diffusers format")
                if needs_prefix:
                    print(f"   - Adding 'transformer.' prefix to keys")
                
                converted_state = {}
                converted_count = 0
                
                for key, value in state_dict.items():
                    new_key = key
                    
                    # 步骤 1: 转换 PEFT 格式到 diffusers 格式
                    if 'lora_A.weight' in new_key:
                        new_key = new_key.replace('lora_A.weight', 'lora.down.weight')
                        converted_count += 1
                    elif 'lora_B.weight' in new_key:
                        new_key = new_key.replace('lora_B.weight', 'lora.up.weight')
                        converted_count += 1
                    
                    # 步骤 2: 添加 transformer 前缀（如果还没有的话）
                    if not new_key.startswith('transformer.'):
                        new_key = f'transformer.{new_key}'
                    
                    converted_state[new_key] = value
                
                if needs_format_conversion:
                    print(f"   ✅ Converted {converted_count} PEFT keys")
                print(f"   ✅ Total keys: {len(converted_state)}")
                
                with tempfile.TemporaryDirectory() as temp_dir:
                    temp_file = os.path.join(temp_dir, weight_name)
                    save_file(converted_state, temp_file)
                    return _original_load_lora_weights(self, temp_dir, **kwargs)
            else:
                print(f"✅ LoRA already in correct format: {lora_file}")
    
    # 不需要转换，使用原始方法
    return _original_load_lora_weights(self, pretrained_model_name_or_path_or_dict, **kwargs)

# 应用 monkey patch
FluxKontextControlPipeline.load_lora_weights = _patched_load_lora_weights
print("✅ Monkey patch applied to FluxKontextPipeline.load_lora_weights")

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
sys.path.append(os.path.abspath(os.path.join(current_dir, '..')))
sys.path.append(os.path.abspath(os.path.join(current_dir, '..', '..', 'comfy_extras')))

from train.src.condition.edge_extraction import InformativeDetector, HEDDetector
from utils_node import BlendInpaint, JoinImageWithAlpha, GrowMask, InvertMask, ColorDetector

TEST_MODE = False

class KontextEditModel():
    def __init__(self, base_model_path="black-forest-labs/FLUX.1-Kontext-dev", device="cuda",
                 aux_lora_dir="models/v2_ckpt", easycontrol_base_dir="models/v2_ckpt",
                 aux_lora_weight_name="puzzle_lora.safetensors",
                 aux_lora_weight=1.0):
        # Keep necessary preprocessors
        self.mask_processor = GrowMask()
        self.scribble_processor = HEDDetector.from_pretrained()
        self.lineart_processor = InformativeDetector.from_pretrained()
        self.color_processor = ColorDetector()
        self.blender = BlendInpaint()

        # Initialize the new pipeline (Kontext version)
        self.device = device
        self.pipe = FluxKontextControlPipeline.from_pretrained(base_model_path, torch_dtype=torch.bfloat16)
        transformer = FluxTransformer2DModel.from_pretrained(
            base_model_path, 
            subfolder="transformer",
            torch_dtype=torch.bfloat16, 
            device=self.device
        )
        self.pipe.transformer = transformer
        self.pipe.to(self.device, dtype=torch.bfloat16)

        control_lora_config = {
            "local": {
                "path": os.path.join(easycontrol_base_dir, "local_lora.safetensors"),
                "lora_weights": [1.0],
                "cond_size": 512,
            },
            "removal": {
                "path": os.path.join(easycontrol_base_dir, "removal_lora.safetensors"),
                "lora_weights": [1.0],
                "cond_size": 512,
            },
            "edge": {
                "path": os.path.join(easycontrol_base_dir, "edge_lora.safetensors"),
                "lora_weights": [1.0],
                "cond_size": 512,
            },
            "color": {
                "path": os.path.join(easycontrol_base_dir, "color_lora.safetensors"),
                "lora_weights": [1.0],
                "cond_size": 512,
            },
        }
        self.pipe.load_control_loras(control_lora_config)

        # Aux LoRA for foreground mode
        self.aux_lora_weight_name = aux_lora_weight_name
        self.aux_lora_dir = aux_lora_dir
        self.aux_lora_weight = aux_lora_weight
        self.aux_adapter_name = "aux"
        
        from safetensors.torch import load_file as _sft_load
        aux_path = os.path.join(self.aux_lora_dir, self.aux_lora_weight_name)
        if os.path.isfile(aux_path):
            self.pipe.load_lora_weights(aux_path, adapter_name=self.aux_adapter_name)
            print(f"Loaded aux LoRA: {aux_path}")
            # Ensure aux LoRA is disabled by default; it will be enabled only in foreground_edit
            self._disable_aux_lora()
        else:
            print(f"Aux LoRA not found at {aux_path}, foreground mode will run without it.")


    # gamma is now applied inside the pipeline based on control_dict

    def _tensor_to_pil(self, tensor_image):
        # Converts a ComfyUI-style tensor [1, H, W, 3] to a PIL Image
        return Image.fromarray(np.clip(255. * tensor_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

    def _pil_to_tensor(self, pil_image):
        # Converts a PIL image to a ComfyUI-style tensor [1, H, W, 3]
        return torch.from_numpy(np.array(pil_image).astype(np.float32) / 255.0).unsqueeze(0)

    def clear_cache(self):
        for name, attn_processor in self.pipe.transformer.attn_processors.items():
            if hasattr(attn_processor, 'bank_kv'):
                attn_processor.bank_kv.clear()
            if hasattr(attn_processor, 'bank_attn'):
                attn_processor.bank_attn = None

    def _enable_aux_lora(self):
        self.pipe.enable_lora()
        self.pipe.set_adapters([self.aux_adapter_name], adapter_weights=[self.aux_lora_weight])
        print(f"Enabled aux LoRA '{self.aux_adapter_name}' with weight {self.aux_lora_weight}")

    def _disable_aux_lora(self):
        self.pipe.disable_lora()
        print("Disabled aux LoRA")

    def _expand_mask(self, mask_tensor: torch.Tensor, expand: int = 0) -> torch.Tensor:
        if expand <= 0:
            return mask_tensor
        expanded = self.mask_processor.expand_mask(mask_tensor, expand=expand, tapered_corners=True)[0]
        return expanded

    def _tensor_mask_to_pil3(self, mask_tensor: torch.Tensor) -> Image.Image:
        mask_01 = torch.clamp(mask_tensor, 0.0, 1.0)
        if mask_01.ndim == 3 and mask_01.shape[-1] == 3:
            mask_01 = mask_01[..., 0]
        if mask_01.ndim == 3 and mask_01.shape[0] == 1:
            mask_01 = mask_01[0]
        pil = self._tensor_to_pil(mask_01.unsqueeze(-1).repeat(1, 1, 3))
        return pil

    def _apply_black_mask(self, image_tensor: torch.Tensor, binary_mask: torch.Tensor) -> Image.Image:
        # image_tensor: [1, H, W, 3] in [0,1]
        # binary_mask: [H, W] or [1, H, W], 1=mask area (white)
        if binary_mask.ndim == 3:
            binary_mask = binary_mask[0]
        mask_bool = (binary_mask > 0.5)
        img = image_tensor.clone()
        img[0][mask_bool] = 0.0
        return self._tensor_to_pil(img)

    def edge_edit(self,
                image, colored_image, positive_prompt, 
                base_mask, add_mask, remove_mask, 
                fine_edge, 
                edge_strength, color_strength,
                seed, steps, cfg):

        generator = torch.Generator(device=self.device).manual_seed(seed)
        
        # Prepare mask and original image
        original_image_tensor = image.clone()
        original_mask = base_mask
        original_mask = self._expand_mask(original_mask, expand=25)
        
        image_pil = self._tensor_to_pil(image)
        # image_pil.save("image_pil.png")
        control_dict = {}
        lineart_output = None

        # Determine control type: color or edge
        if not torch.equal(image, colored_image):
            print("Apply color control")
            colored_image_pil = self._tensor_to_pil(colored_image)
            # Create color block condition
            color_image_np = np.array(colored_image_pil)
            downsampled = cv2.resize(color_image_np, (32, 32), interpolation=cv2.INTER_AREA)
            upsampled = cv2.resize(downsampled, (256, 256), interpolation=cv2.INTER_NEAREST)
            color_block = Image.fromarray(upsampled)
            # Create grayscale condition
            
            control_dict = {
                "type": "color",
                "spatial_images": [color_block],
                "gammas": [color_strength]
            }
        else:
            print("Apply edge control")
            if fine_edge == "enable":
                lineart_image = self.lineart_processor(np.array(self._tensor_to_pil(image.cpu().squeeze())), detect_resolution=1024, style="contour", output_type="pil")
                lineart_output = self._pil_to_tensor(lineart_image)
            else:
                scribble_image = self.scribble_processor(np.array(self._tensor_to_pil(image.cpu().squeeze())), safe=True, resolution=512, output_type="pil")
                lineart_output = self._pil_to_tensor(scribble_image)
            
            if lineart_output is None:
                raise ValueError("Preprocessor failed to generate lineart.")

            # Apply user sketches to the lineart
            add_mask_resized = F.interpolate(add_mask.unsqueeze(0).float(), size=(lineart_output.shape[1], lineart_output.shape[2]), mode='nearest').squeeze(0)
            remove_mask_resized = F.interpolate(remove_mask.unsqueeze(0).float(), size=(lineart_output.shape[1], lineart_output.shape[2]), mode='nearest').squeeze(0)

            bool_add_mask_resized = (add_mask_resized > 0.5)
            bool_remove_mask_resized = (remove_mask_resized > 0.5)

            lineart_output[bool_remove_mask_resized] = 0.0
            lineart_output[bool_add_mask_resized] = 1.0

            control_dict = {
                "type": "edge",
                "spatial_images": [self._tensor_to_pil(lineart_output)],
                "gammas": [edge_strength]
            }

        # Prepare debug/output images
        colored_image_np = np.array(self._tensor_to_pil(colored_image))
        debug_image = lineart_output if lineart_output is not None else self.color_processor(colored_image_np, detect_resolution=1024, output_type="pil")

        # Run inference
        result_pil = self.pipe(
            prompt=positive_prompt,
            image=image_pil,
            height=image_pil.height,
            width=image_pil.width,
            guidance_scale=cfg,
            num_inference_steps=steps,
            generator=generator,
            max_sequence_length=128,
            control_dict=control_dict,
        ).images[0]

        self.clear_cache()

        # result_pil.save("result_pil.png")
        result_tensor = self._pil_to_tensor(result_pil)
        # final_image = self.blender.blend_inpaint(result_tensor, original_image_tensor, original_mask, kernel=10, sigma=10)[0]
        final_image = result_tensor
        return (final_image, debug_image, original_mask)

    def object_removal(self,
                       image, positive_prompt, 
                       remove_mask, 
                       local_strength,
                       seed, steps, cfg):
        
        generator = torch.Generator(device=self.device).manual_seed(seed)

        original_image_tensor = image.clone()
        original_mask = remove_mask
        original_mask = self._expand_mask(remove_mask, expand=10)
        
        image_pil = self._tensor_to_pil(image)
        # image_pil.save("image_pil.png")
        # Prepare spatial image: original masked to black in the remove area
        spatial_pil = self._apply_black_mask(image, original_mask)
        # spatial_pil.save("spatial_pil.png")
        # Note: mask is not passed to pipeline; we use it only for blending
        control_dict = {
            "type": "removal",
            "spatial_images": [spatial_pil],
            "gammas": [local_strength]
        }

        result_pil = self.pipe(
            prompt=positive_prompt,
            image=image_pil,
            height=image_pil.height,
            width=image_pil.width,
            guidance_scale=cfg,
            num_inference_steps=steps,
            generator=generator,
            control_dict=control_dict,
        ).images[0]

        self.clear_cache()

        result_tensor = self._pil_to_tensor(result_pil)
        final_image = self.blender.blend_inpaint(result_tensor, original_image_tensor, original_mask, kernel=10, sigma=10)[0]
        # final_image = result_tensor
        return (final_image, self._pil_to_tensor(spatial_pil), original_mask)

    def local_edit(self,
                   image, positive_prompt, fill_mask, local_strength,
                   seed, steps, cfg):
        generator = torch.Generator(device=self.device).manual_seed(seed)
        original_image_tensor = image.clone()
        original_mask = self._expand_mask(fill_mask, expand=10)
        image_pil = self._tensor_to_pil(image)
        # image_pil.save("image_pil.png")

        spatial_pil = self._apply_black_mask(image, original_mask)
        # spatial_pil.save("spatial_pil.png")
        control_dict = {
            "type": "local",
            "spatial_images": [spatial_pil],
            "gammas": [local_strength]
        }

        result_pil = self.pipe(
            prompt=positive_prompt,
            image=image_pil,
            height=image_pil.height,
            width=image_pil.width,
            guidance_scale=cfg,
            num_inference_steps=steps,
            generator=generator,
            max_sequence_length=128,
            control_dict=control_dict,
        ).images[0]

        self.clear_cache()
        result_tensor = self._pil_to_tensor(result_pil)
        final_image = self.blender.blend_inpaint(result_tensor, original_image_tensor, original_mask, kernel=10, sigma=10)[0]
        # final_image = result_tensor
        return (final_image, self._pil_to_tensor(spatial_pil), original_mask)

    def foreground_edit(self,
                        merged_image, positive_prompt,
                        add_prop_mask, fill_mask, fix_perspective, grow_size,
                        seed, steps, cfg):
        generator = torch.Generator(device=self.device).manual_seed(seed)

        edit_mask = torch.clamp(self._expand_mask(add_prop_mask, expand=grow_size) + fill_mask, 0.0, 1.0)
        final_mask = self._expand_mask(edit_mask, expand=25)
        if fix_perspective == "enable":
            positive_prompt = positive_prompt + " Fix the perspective if necessary."
        # Prepare edited input image: inside edit_mask but outside add_prop_mask set to white
        img = merged_image.clone()
        base_mask = (edit_mask > 0.5)
        add_only = (add_prop_mask <= 0.5) & base_mask  # [1, H, W] bool
        add_only_3 = add_only.squeeze(0).unsqueeze(-1).expand(-1, -1, img.shape[-1])  # [H, W, 3]
        img[0] = torch.where(add_only_3, torch.ones_like(img[0]), img[0])

        image_pil = self._tensor_to_pil(img)
        # image_pil.save("image_pil.png")

        # Enable aux LoRA only for foreground
        self._enable_aux_lora()

        result_pil = self.pipe(
            prompt=positive_prompt,
            image=image_pil,
            height=image_pil.height,
            width=image_pil.width,
            guidance_scale=cfg,
            num_inference_steps=steps,
            generator=generator,
            max_sequence_length=128,
            control_dict=None,
        ).images[0]

        # Disable aux LoRA afterwards
        self._disable_aux_lora()

        self.clear_cache()
        final_image = self._pil_to_tensor(result_pil)
        # final_image = self.blender.blend_inpaint(final_image, img, final_mask, kernel=10, sigma=10)[0]
        return (final_image, self._pil_to_tensor(image_pil), edit_mask)

    def kontext_edit(self,
                     image, positive_prompt,
                     seed, steps, cfg):
        generator = torch.Generator(device=self.device).manual_seed(seed)
        image_pil = self._tensor_to_pil(image)

        result_pil = self.pipe(
            prompt=positive_prompt,
            image=image_pil,
            height=image_pil.height,
            width=image_pil.width,
            guidance_scale=cfg,
            num_inference_steps=steps,
            generator=generator,
            max_sequence_length=128,
            control_dict=None,
        ).images[0]

        final_image = self._pil_to_tensor(result_pil)
        mask = torch.zeros((1, final_image.shape[1], final_image.shape[2]), dtype=torch.float32, device=final_image.device)
        return (final_image, image, mask)

    def process(self, image, colored_image, 
                 merged_image, positive_prompt,
                total_mask, add_mask, remove_mask, add_prop_mask, fill_mask, 
                fine_edge, fix_perspective, edge_strength, color_strength, local_strength, grow_size,
                seed, steps, cfg, flag="precise_edit"):
        if flag == "foreground":
            return self.foreground_edit(merged_image, positive_prompt, add_prop_mask, fill_mask, fix_perspective, grow_size, seed, steps, cfg)
        elif flag == "local":
            return self.local_edit(image, positive_prompt, fill_mask, local_strength, seed, steps, cfg)
        elif flag == "removal":
            return self.object_removal(image, positive_prompt, remove_mask, local_strength, seed, steps, cfg)
        elif flag == "precise_edit":
            return self.edge_edit(
                image, colored_image, positive_prompt,
                total_mask, add_mask, remove_mask,
                fine_edge,
                edge_strength, color_strength,
                seed, steps, cfg
            )
        elif flag == "kontext":
            return self.kontext_edit(image, positive_prompt, seed, steps, cfg)
        else:
            raise ValueError("Invalid Editing Type: {}".format(flag))