utils.py

"""
Utility functions for Qwen Image Edit ComfyUI plugin
"""

import torch
from PIL import Image
import numpy as np
from typing import Union, List, Tuple, Optional


def tensor_to_pil(tensor: torch.Tensor, batch_idx: int = 0) -> Image.Image:
    """Convert ComfyUI tensor to PIL Image.
    
    Args:
        tensor: ComfyUI IMAGE tensor (BCHW format, float32 [0,1])
        batch_idx: Which image in batch to convert
        
    Returns:
        PIL Image in RGB format
    """
    if tensor.dim() == 4:
        tensor = tensor[batch_idx]
    elif tensor.dim() == 3:
        pass  # already single image CHW
    else:
        raise ValueError(f"Expected 3D or 4D tensor, got {tensor.dim()}D")
    
    # Convert from CHW to HWC and scale to [0,255]
    image_np = (tensor.clamp(0, 1).permute(1, 2, 0).detach().cpu().numpy() * 255).astype(np.uint8)
    return Image.fromarray(image_np)


def pil_to_tensor(image: Image.Image) -> torch.Tensor:
    """Convert PIL Image to ComfyUI tensor format.
    
    Args:
        image: PIL Image in RGB format
        
    Returns:
        ComfyUI IMAGE tensor (1CHW format, float32 [0,1])
    """
    # Convert to numpy array and normalize to [0,1]
    arr = np.array(image).astype(np.float32) / 255.0
    
    # Handle grayscale by converting to RGB
    if arr.ndim == 2:
        arr = np.stack([arr, arr, arr], axis=-1)
    
    # Remove alpha channel if present
    if arr.shape[-1] == 4:
        arr = arr[..., :3]
    
    # Convert from HWC to 1CHW
    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).contiguous()
    return tensor


def batch_tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
    """Convert batch tensor to list of PIL Images.
    
    Args:
        tensor: ComfyUI IMAGE tensor (BCHW format)
        
    Returns:
        List of PIL Images
    """
    if tensor.dim() != 4:
        raise ValueError(f"Expected 4D batch tensor, got {tensor.dim()}D")
    
    return [tensor_to_pil(tensor, i) for i in range(tensor.shape[0])]


def pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
    """Convert list of PIL Images to batch tensor.
    
    Args:
        images: List of PIL Images
        
    Returns:
        ComfyUI IMAGE tensor (BCHW format)
    """
    if not images:
        raise ValueError("Empty image list")
    
    tensors = [pil_to_tensor(img) for img in images]
    return torch.cat(tensors, dim=0)


def validate_gpu_memory() -> Tuple[bool, str]:
    """Check if sufficient GPU memory is available.
    
    Returns:
        Tuple of (is_sufficient, message)
    """
    if not torch.cuda.is_available():
        return False, "CUDA not available. CPU mode will be used but may be very slow."
    
    try:
        memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        if memory_gb < 8:
            return False, f"Only {memory_gb:.1f}GB GPU memory available. 8GB+ recommended for optimal performance."
        return True, f"GPU memory: {memory_gb:.1f}GB - sufficient for Qwen Image Edit"
    except Exception as e:
        return False, f"Could not check GPU memory: {e}"


def safe_model_load(model_path: str, device: str = "auto") -> Tuple[bool, str]:
    """Safely attempt to load model with appropriate device selection.
    
    Args:
        model_path: Path or name of the model
        device: Target device ("auto", "cuda", "cpu")
        
    Returns:
        Tuple of (success, message/error)
    """
    if device == "auto":
        if torch.cuda.is_available():
            gpu_ok, gpu_msg = validate_gpu_memory()
            device = "cuda" if gpu_ok else "cpu"
            if not gpu_ok:
                return False, f"Falling back to CPU: {gpu_msg}"
        else:
            device = "cpu"
    
    try:
        # This is just a validation function - actual loading happens in the node
        return True, f"Model loading prepared for device: {device}"
    except Exception as e:
        return False, f"Model loading failed: {e}"


def prepare_generation_params(params: Optional[dict], defaults: Optional[dict] = None) -> dict:
    """Prepare and validate generation parameters.
    
    Args:
        params: User provided parameters
        defaults: Default parameter values
        
    Returns:
        Merged and validated parameters
    """
    if defaults is None:
        defaults = {
            "true_cfg_scale": 4.0,
            "negative_prompt": " ",
            "num_inference_steps": 50,
            "seed": 0,
        }
    
    if params is None:
        params = {}
    
    # Merge with defaults
    result = defaults.copy()
    result.update(params)
    
    # Validate ranges
    result["true_cfg_scale"] = max(0.0, min(20.0, float(result["true_cfg_scale"])))
    result["num_inference_steps"] = max(1, min(200, int(result["num_inference_steps"])))
    result["seed"] = max(0, min(2**31-1, int(result["seed"])))
    
    return result


def split_prompts(prompt_text: str) -> List[str]:
    """Split multi-line prompt text into individual prompts.
    
    Args:
        prompt_text: Multi-line string with prompts
        
    Returns:
        List of individual prompt strings
    """
    lines = [line.strip() for line in prompt_text.splitlines() if line.strip()]
    return lines if lines else [""]


def format_error_message(error: Exception, context: str = "") -> str:
    """Format error messages for user-friendly display.
    
    Args:
        error: The exception that occurred
        context: Additional context about where the error occurred
        
    Returns:
        Formatted error message
    """
    error_type = type(error).__name__
    error_msg = str(error)
    
    if context:
        return f"[{context}] {error_type}: {error_msg}"
    else:
        return f"{error_type}: {error_msg}"