ZeST

Runtime error

App Files Files Community

ZeST / run_batch.py

fffiloni

Upload 47 files

a9289c0 about 2 months ago

raw history blame contribute delete

No virus

3.86 kB

	from diffusers import StableDiffusionXLControlNetInpaintPipeline, ControlNetModel
	from rembg import remove
	from PIL import Image
	import torch
	from ip_adapter import IPAdapterXL
	from ip_adapter.utils import register_cross_attention_hook, get_net_attn_map, attnmaps2images
	from PIL import Image, ImageChops
	from PIL import ImageEnhance
	import numpy as np
	import glob

	def image_grid(imgs, rows, cols):
	assert len(imgs) == rows*cols

	w, h = imgs[0].size
	grid = Image.new('RGB', size=(colsw, rowsh))
	grid_w, grid_h = grid.size

	for i, img in enumerate(imgs):
	grid.paste(img, box=(i%colsw, i//colsh))
	return grid

	base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
	image_encoder_path = "models/image_encoder"
	ip_ckpt = "sdxl_models/ip-adapter_sdxl_vit-h.bin"
	controlnet_path = "diffusers/controlnet-depth-sdxl-1.0"
	device = "cuda"

	torch.cuda.empty_cache()

	# load SDXL pipeline
	controlnet = ControlNetModel.from_pretrained(controlnet_path, variant="fp16", use_safetensors=True, torch_dtype=torch.float16).to(device)
	pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
	base_model_path,
	controlnet=controlnet,
	use_safetensors=True,
	torch_dtype=torch.float16,
	add_watermarker=False,
	).to(device)
	pipe.unet = register_cross_attention_hook(pipe.unet)

	ip_model = IPAdapterXL(pipe, image_encoder_path, ip_ckpt, device)



	textures = [tex.split('/')[-1].replace('.png', '') for tex in glob.glob('demo_assets/material_exemplars/*.png')]
	objs = [obj.split('/')[-1].replace('.png', '') for obj in glob.glob('demo_assets/input_imgs/*.png')]

	for texture in textures:
	for obj in objs:
	target_image_path = 'demo_assets/input_imgs/' + obj + '.png' # Replace with your image path
	target_image = Image.open(target_image_path).convert('RGB')
	rm_bg = remove(target_image)
	# output.save(output_path)
	target_mask = rm_bg.convert("RGB").point(lambda x: 0 if x < 1 else 255).convert('L').convert('RGB')# Convert mask to grayscale

	# Ensure mask is the same size as image

	# mask = ImageChops.invert(mask)
	# Generate random noise for the size of the image
	noise = np.random.randint(0, 256, target_image.size + (3,), dtype=np.uint8)
	noise_image = Image.fromarray(noise)
	mask_target_img = ImageChops.lighter(target_image, target_mask)
	invert_target_mask = ImageChops.invert(target_mask)


	gray_target_image = target_image.convert('L').convert('RGB')
	gray_target_image = ImageEnhance.Brightness(gray_target_image)

	# Adjust brightness
	# The factor 1.0 means original brightness, greater than 1.0 makes the image brighter. Adjust this if the image is too dim
	factor = 1.0 # Try adjusting this to get the desired brightness

	gray_target_image = gray_target_image.enhance(factor)
	grayscale_img = ImageChops.darker(gray_target_image, target_mask)
	img_black_mask = ImageChops.darker(target_image, invert_target_mask)
	grayscale_init_img = ImageChops.lighter(img_black_mask, grayscale_img)
	init_img = grayscale_init_img

	ip_image = Image.open("demo_assets/material_exemplars/" + texture + ".png")
	np_image = np.array(Image.open('demo_assets/depths/' + obj + '.png'))

	np_image = (np_image / 256).astype('uint8')

	depth_map = Image.fromarray(np_image).resize((1024,1024))

	init_img = init_img.resize((1024,1024))
	mask = target_mask.resize((1024, 1024))

	num_samples = 1
	images = ip_model.generate(pil_image=ip_image, image=init_img, control_image=depth_map, mask_image=mask, controlnet_conditioning_scale=0.9, num_samples=num_samples, num_inference_steps=30, seed=42)
	images[0].save('demo_assets/output_images/' + obj + '_' + texture + '.png' )