Spaces:
Runtime error
Runtime error
from logger import rich_logger as l | |
from wandb.integration.diffusers import autolog | |
from config import Project_Name | |
from clear_memory import clear_memory | |
from typing import List | |
import numpy as np | |
import torch | |
from PIL import Image | |
from mask_generator import convert_to_numpy_array, generate_mask | |
from diffusers.utils import load_image | |
import cv2 | |
from config import controlnet_adapter_model_name,controlnet_base_model_name | |
from diffusers import ControlNetModel,StableDiffusionControlNetInpaintPipeline | |
autolog(init=dict(project=Project_Name)) | |
def make_inpaint_condition(init_image, mask_image): | |
# Prepare control image | |
init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0 | |
mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0 | |
assert init_image.shape[0:1] == mask_image.shape[0:1], "image and image_mask must have the same image size" | |
init_image[mask_image > 0.5] = -1.0 # set as masked pixel | |
init_image = np.expand_dims(init_image, 0).transpose(0, 3, 1, 2) | |
init_image = torch.from_numpy(init_image) | |
return init_image | |
def make_image_controlnet(image, | |
mask_image, | |
controlnet_conditioning_image, | |
positive_prompt: str, negative_prompt: str, | |
seed: int = 2356132) -> List[Image.Image]: | |
"""Method to make image using controlnet | |
Args: | |
image (np.ndarray): input image | |
mask_image (np.ndarray): mask image | |
controlnet_conditioning_image (np.ndarray): conditioning image | |
positive_prompt (str): positive prompt string | |
negative_prompt (str): negative prompt string | |
seed (int, optional): seed. Defaults to 2356132. | |
Returns: | |
List[Image.Image]: list of generated images | |
""" | |
controlnet = ControlNetModel.from_pretrained(controlnet_adapter_model_name, torch_dtype=torch.float32) | |
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( | |
controlnet_base_model_name, controlnet=controlnet, torch_dtype=torch.float32 | |
) | |
image = pipe(prompt=positive_prompt,negative_prompt=negative_prompt, image=init_image, mask_image=mask_image, control_image=controlnet_conditioning_image).images[0] | |
return image | |
if __name__ == "__main__": | |
init_image = load_image('/home/product_diffusion_api/sample_data/example1.jpg') | |
mask_image = load_image('/home/product_diffusion_api/scripts/mask.jpg') | |
controlnet_conditioning_image = make_inpaint_condition(init_image=init_image,mask_image=mask_image) | |
result = make_image_controlnet(positive_prompt="Product used in kitchen 4k natural photography",negative_prompt="No artifcats",image=init_image,mask_image=mask_image,controlnet_conditioning_image=controlnet_conditioning_image) | |