Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,260 Bytes
5a0778e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import torch
from diffusers.pipelines.flux.pipeline_flux_fill_unmasked_image_condition_version import FluxFillPipeline_token12_depth_only as FluxFillPipeline
from diffusers.utils import load_image
import os, glob
import numpy as np
import cv2
from PIL import Image, ImageOps
image_path = ["example_data/I-210618_I01001_W01_I-210618_I01001_W01_F0153_img.jpg"]
pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights("stage1/checkpoint-4800")
for image_ep in image_path:
mask_path = image_ep.replace("_img.jpg","_mask.png")
image = Image.open(image_ep) # place_hold
depth = Image.open(image_ep.replace("_img.jpg",
"_depth_img.png"))
image_name = os.path.basename(image_ep)
mask = Image.open(mask_path).convert("L")
mask = ImageOps.invert(mask) # inverse rord_mask
# mask_np = np.array(mask)
# # mask dilation
# dilation_px = 32
# kernel = np.ones((3, 3), np.uint8)
# iterations = dilation_px // 2
# dilated_mask = cv2.dilate(mask_np, kernel, iterations=iterations)
# mask = Image.fromarray(dilated_mask)
orig_w, orig_h = image.size
# Resize to 1024 × 1024
# target_size = (1024, 1024)
# image_resized = image.resize(target_size, Image.BICUBIC)
# mask_resized = mask.resize(target_size, Image.NEAREST)
# depth_resized = depth.resize(target_size, Image.BICUBIC)
w, h = image.size
MAX_SIZE = 1024
if max(w, h) > MAX_SIZE:
factor = MAX_SIZE / max(w, h)
w = int(factor * w)
h = int(factor * h)
width, height = map(lambda x: x - x % 64, (w, h))
image_out = pipe(
prompt="A beautiful scene",
image=image,
mask_image=mask,
width=width,
height=height,
guidance_scale=30,
num_inference_steps=50,
max_sequence_length=512,
generator=torch.Generator("cpu").manual_seed(0),
depth=depth
).images[0]
image_final = image_out.resize((orig_w, orig_h), Image.BICUBIC)
output_dir = "./depth_fillin_results"
os.makedirs(output_dir, exist_ok=True)
image_final.save(os.path.join(output_dir, image_name))
|