'''
import os
from datasets import load_dataset
from PIL import Image

def save_images_locally(dataset_name="svjack/InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV_text_mask"):
    # 加载数据集
    dataset = load_dataset(dataset_name)["train"]

    # 创建文件夹
    image_dir = "original_images"
    mask_dir = "mask_images"
    os.makedirs(image_dir, exist_ok=True)
    os.makedirs(mask_dir, exist_ok=True)

    # 遍历数据集并保存图像
    for idx, example in enumerate(dataset):
        # 保存原图
        image_path = os.path.join(image_dir, f"{idx:04d}.png")
        example["Wang_Leehom_poster_image"].save(image_path)

        # 保存 mask 图像
        mask_path = os.path.join(mask_dir, f"{idx:04d}.png")
        example["mask_image"].convert("RGB").save(mask_path)

    print(f"✅ 原图已保存至 {image_dir}")
    print(f"✅ Mask 图像已保存至 {mask_dir}")

if __name__ == "__main__":
    save_images_locally()

import cv2
import numpy as np
from PIL import Image

def blur_mask_opencv(
    mask: Image.Image, 
    blur_factor: int, 
    edge_expand: int = 0  # 新增参数：边缘扩展像素值
) -> Image.Image:
    """
    增强版遮罩处理：支持边缘扩展 + 高斯模糊羽化
    参数:
        mask: PIL格式的遮罩图像（支持 RGBA/RGB/L 模式）
        blur_factor: 模糊强度（奇数）
        edge_expand: 边缘扩展像素值（非负整数）
    """
    # 转换为OpenCV格式（NumPy数组）
    mask_np = np.array(mask)
    original_mode = mask.mode  # 保存原始色彩模式
    
    # 通道分离逻辑
    if original_mode == 'RGBA':
        alpha = mask_np[:, :, 3]
        processed = alpha
    elif original_mode == 'RGB':
        processed = cv2.cvtColor(mask_np, cv2.COLOR_RGB2GRAY)
    else:  # L模式（灰度）
        processed = mask_np
    
    # 边缘扩展（核心新增功能）
    if edge_expand > 0:
        kernel = np.ones((3, 3), np.uint8)  # 3x3方形结构核
        processed = cv2.dilate(
            processed, 
            kernel, 
            iterations=edge_expand  # 扩展次数 = 扩展像素值
        )
    
    # 高斯模糊羽化
    kernel_size = blur_factor | 1  # 确保为奇数
    blurred = cv2.GaussianBlur(processed, (kernel_size, kernel_size), 0)
    
    # 重建为PIL图像
    if original_mode == 'RGBA':
        result = mask_np.copy()
        result[:, :, 3] = blurred  # 仅替换Alpha通道
        return Image.fromarray(result)
    else:
        return Image.fromarray(blurred)


#from PIL import Image 
#mask = Image.open("image (51).jpg")

# 使用示例
#blurred_mask = blur_mask_opencv(mask, 33, 20)
#blurred_mask.save("b_mask.png")
#blurred_mask


import os
import numpy as np
from PIL import Image
import torch
from diffusers import FluxFillPipeline

# 初始化模型
pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()

def calculate_optimal_dimensions(image):
    original_width, original_height = image.size
    MIN_ASPECT_RATIO = 9 / 16
    MAX_ASPECT_RATIO = 16 / 9
    FIXED_DIMENSION = 1024
    original_aspect_ratio = original_width / original_height

    if original_aspect_ratio > 1:
        width = FIXED_DIMENSION
        height = round(FIXED_DIMENSION / original_aspect_ratio)
    else:
        height = FIXED_DIMENSION
        width = round(FIXED_DIMENSION * original_aspect_ratio)

    width = (width // 8) * 8
    height = (height // 8) * 8
    calculated_aspect_ratio = width / height

    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
        width = (height * MAX_ASPECT_RATIO // 8) * 8
    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
        height = (width / MIN_ASPECT_RATIO // 8) * 8

    width = max(width, 576) if width == FIXED_DIMENSION else width
    height = max(height, 576) if height == FIXED_DIMENSION else height
    return width, height

def calculate_white_pixel_ratio(mask_image):
    mask_array = np.array(mask_image.convert('L'))
    white_pixels = np.sum(mask_array == 255)
    total_pixels = mask_array.size
    return white_pixels / total_pixels

def inpaint(image_path, mask_path, output_path):
    image = Image.open(image_path).convert("RGB")
    mask = Image.open(mask_path).convert("L")
    #mask = blur_mask_opencv(mask, 33, 20)
    mask = blur_mask_opencv(mask, 33, 60)

    white_ratio = calculate_white_pixel_ratio(mask)
    if white_ratio >= 1 / 3:
        image.save(output_path)
        return

    width, height = calculate_optimal_dimensions(image)
    result = pipe(
        prompt="",
        height=height,
        width=width,
        image=image,
        mask_image=mask,
        num_inference_steps=40,
        guidance_scale=28,
    ).images[0]
    result.save(output_path)

def process_all_image_pairs():
    image_dir = "original_images"
    mask_dir = "mask_images"
    output_dir = "inpaint_results"
    os.makedirs(output_dir, exist_ok=True)

    for filename in os.listdir(image_dir):
        base_name = os.path.splitext(filename)[0]
        image_path = os.path.join(image_dir, filename)
        mask_path = os.path.join(mask_dir, filename)
        output_path = os.path.join(output_dir, filename)

        if not os.path.exists(mask_path):
            continue

        try:
            inpaint(image_path, mask_path, output_path)
        except Exception as e:
            print(f"❌ 处理 {filename} 时出错: {e}")

    print(f"✅ 修复结果已保存至 {output_dir}")

if __name__ == "__main__":
    process_all_image_pairs()

import os
from datasets import load_dataset, DatasetDict, Image as HfImage
from PIL import Image

def update_dataset_with_inpaint_results():
    dataset_name = "svjack/InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV_text_mask"
    input_dataset = load_dataset(dataset_name)["train"]

    output_dir = "inpaint_results"
    output_dataset_path = "InfiniteYou_PosterCraft_Wang_Leehom_Poster_FP8_WAV_text_mask_inpaint"
    os.makedirs(output_dataset_path, exist_ok=True)

    def add_inpaint_image(example, idx):
        output_path = os.path.join(output_dir, f"{idx:04d}.png")
        
        if os.path.exists(output_path):
            example["inpaint_image"] = output_path
        else:
            # 如果没有生成修复图像，则使用原始图像
            example["inpaint_image"] = example["Wang_Leehom_poster_image"]
        
        return example

    updated_dataset = input_dataset.map(
        lambda ex, idx: add_inpaint_image(ex, idx),
        with_indices=True,
        batched=False,
        num_proc=1
    )

    updated_dataset = updated_dataset.cast_column("inpaint_image", HfImage())
    updated_dataset.save_to_disk(output_dataset_path)

    print(f"✅ 更新后的数据集已保存至 {output_dataset_path}")

if __name__ == "__main__":
    update_dataset_with_inpaint_results()
'''

import torch
import spaces
import gradio as gr
from diffusers import FluxFillPipeline

pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")

# reference https://huggingface.co/spaces/black-forest-labs/FLUX.1-Fill-dev/blob/main/app.py 
def calculate_optimal_dimensions(image):
    # Extract the original dimensions
    original_width, original_height = image.size
    
    # Set constants
    MIN_ASPECT_RATIO = 9 / 16
    MAX_ASPECT_RATIO = 16 / 9
    FIXED_DIMENSION = 1024

    # Calculate the aspect ratio of the original image
    original_aspect_ratio = original_width / original_height

    # Determine which dimension to fix
    if original_aspect_ratio > 1:  # Wider than tall
        width = FIXED_DIMENSION
        height = round(FIXED_DIMENSION / original_aspect_ratio)
    else:  # Taller than wide
        height = FIXED_DIMENSION
        width = round(FIXED_DIMENSION * original_aspect_ratio)

    # Ensure dimensions are multiples of 8
    width = (width // 8) * 8
    height = (height // 8) * 8

    # Enforce aspect ratio limits
    calculated_aspect_ratio = width / height
    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
        width = (height * MAX_ASPECT_RATIO // 8) * 8
    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
        height = (width / MIN_ASPECT_RATIO // 8) * 8

    # Ensure width and height remain above the minimum dimensions
    width = max(width, 576) if width == FIXED_DIMENSION else width
    height = max(height, 576) if height == FIXED_DIMENSION else height

    return width, height


@spaces.GPU(duration=120)
def inpaint(
    image,
    mask,
    prompt="",
    num_inference_steps=28,
    guidance_scale=50,
):
    image = image.convert("RGB")
    mask = mask.convert("L")
    width, height = calculate_optimal_dimensions(image)

    result = pipe(
        prompt=prompt,
        height= height,
        width= width,
        image= image,
        mask_image=mask,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
    ).images[0]

    result = result.convert("RGBA")

    return result


demo = gr.Interface(
    fn=inpaint,
    inputs=[
        gr.Image(label="image", type="pil"),
        gr.Image(label="mask", type="pil"),
        gr.Text(label="prompt"),
        gr.Number(value=40, label="num_inference_steps"),
        gr.Number(value=28, label="guidance_scale"),
    ],
    outputs=["image"],
    api_name="inpaint",
    examples=[["./assets/rocket.png", "./assets/Inpainting mask.png"]],
    cache_examples=False,
    description="it is recommended that you use https://github.com/la-voliere/react-mask-editor when creating an image mask in JS and then inverse it before sending it to this space",
)

#demo.launch(debug=True,show_error=True)
demo.launch(debug=True,show_error=True, share = True)