File size: 3,742 Bytes
f56ede2 |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import cv2
import torch
from PIL import Image
import numpy as np
import yaml
import argparse
from controlnet_aux import OpenposeDetector
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
from utils.download import load_image
from utils.plot import image_grid
def load_config(config_path):
with open(config_path, 'r') as file:
return yaml.safe_load(file)
def initialize_controlnet(config):
model_id = config['model_id']
local_dir = config.get('local_dir', model_id)
return ControlNetModel.from_pretrained(
local_dir if local_dir != model_id else model_id,
torch_dtype=torch.float16
)
def initialize_pipeline(controlnet, config):
model_id = config['model_id']
local_dir = config.get('local_dir', model_id)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
local_dir if local_dir != model_id else model_id,
controlnet=controlnet,
torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
return pipe
def setup_device(pipe):
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
pipe.enable_model_cpu_offload()
pipe.to(device)
return device
def generate_images(pipe, prompts, pose_images, generators, negative_prompts, num_steps):
return pipe(
prompts,
pose_images,
negative_prompt=negative_prompts,
generator=generators,
num_inference_steps=num_steps
).images
def infer(args):
# Load configuration
configs = load_config(args.config_path)
# Initialize models
controlnet_detector = OpenposeDetector.from_pretrained(
configs[2]['model_id'] # lllyasviel/ControlNet
)
controlnet = initialize_controlnet(configs[0])
pipe = initialize_pipeline(controlnet, configs[1])
# Setup device
device = setup_device(pipe)
# Load and process image
demo_image = load_image(args.image_url)
poses = [controlnet_detector(demo_image)]
# Generate images
generators = [torch.Generator(device="cpu").manual_seed(args.seed) for _ in range(len(poses))]
output_images = generate_images(
pipe,
[args.prompt] * len(generators),
poses,
generators,
[args.negative_prompt] * len(generators),
args.num_steps
)
# Display results
# image_grid(output_images, 2, 2)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="ControlNet image generation with pose detection")
parser.add_argument("--config_path", type=str, default="configs/model_ckpts.yaml",
help="Path to configuration YAML file")
parser.add_argument("--image_url", type=str,
default="https://huggingface.co/datasets/YiYiXu/controlnet-testing/resolve/main/yoga1.jpeg",
help="URL of input image")
parser.add_argument("--prompt", type=str, default="a man is doing yoga",
help="Text prompt for image generation")
parser.add_argument("--negative_prompt", type=str,
default="monochrome, lowres, bad anatomy, worst quality, low quality",
help="Negative prompt for image generation")
parser.add_argument("--num_steps", type=int, default=20,
help="Number of inference steps")
parser.add_argument("--seed", type=int, default=2,
help="Random seed for generation")
# return parser.parse_args()
args = parser.parse_args()
infer(args) |