flux-inpainting / app.py
svjack's picture
Update app.py
aa24a87 verified
'''
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)