danhtran2mind commited on
Commit
a9fda10
·
verified ·
1 Parent(s): 73f8596

Delete src/controlnet_image_generator/old2-infer.py

Browse files
src/controlnet_image_generator/old2-infer.py DELETED
@@ -1,189 +0,0 @@
1
- import cv2
2
- import torch
3
- from PIL import Image
4
- import numpy as np
5
- import yaml
6
- import argparse
7
- from controlnet_aux import OpenposeDetector
8
- from diffusers import (
9
- StableDiffusionControlNetPipeline,
10
- ControlNetModel,
11
- UniPCMultistepScheduler
12
- )
13
-
14
- from utils.download import load_image
15
- from utils.plot import image_grid
16
- import os
17
- from tqdm import tqdm
18
- import re
19
- import uuid
20
-
21
- def load_config(config_path):
22
- try:
23
- with open(config_path, 'r') as file:
24
- return yaml.safe_load(file)
25
- except Exception as e:
26
- raise ValueError(f"Error loading config file: {e}")
27
-
28
- def initialize_controlnet(config):
29
- model_id = config['model_id']
30
- local_dir = config.get('local_dir', model_id)
31
- return ControlNetModel.from_pretrained(
32
- local_dir if local_dir != model_id else model_id,
33
- torch_dtype=torch.float16
34
- )
35
-
36
- def initialize_pipeline(controlnet, config):
37
- model_id = config['model_id']
38
- local_dir = config.get('local_dir', model_id)
39
- pipe = StableDiffusionControlNetPipeline.from_pretrained(
40
- local_dir if local_dir != model_id else model_id,
41
- controlnet=controlnet,
42
- torch_dtype=torch.float16
43
- )
44
- pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
45
- return pipe
46
-
47
- def setup_device(pipe):
48
- device = "cuda" if torch.cuda.is_available() else "cpu"
49
- if device == "cuda":
50
- pipe.enable_model_cpu_offload()
51
- pipe.to(device)
52
- return device
53
-
54
- def generate_images(pipe, prompts, pose_images, generators, negative_prompts, num_steps, guidance_scale, controlnet_conditioning_scale, width, height):
55
- return pipe(
56
- prompts,
57
- pose_images,
58
- negative_prompt=negative_prompts,
59
- generator=generators,
60
- num_inference_steps=num_steps,
61
- guidance_scale=guidance_scale,
62
- controlnet_conditioning_scale=controlnet_conditioning_scale,
63
- width=width,
64
- height=height
65
- ).images
66
-
67
- def infer(args):
68
- # Load configuration
69
- configs = load_config(args.config_path)
70
-
71
- # Initialize models
72
- controlnet_detector = OpenposeDetector.from_pretrained(
73
- configs[2]['model_id'] # lllyasviel/ControlNet
74
- )
75
- controlnet = initialize_controlnet(configs[0])
76
- pipe = initialize_pipeline(controlnet, configs[1])
77
-
78
- # Setup device
79
- device = setup_device(pipe)
80
-
81
- # Load and process image
82
- try:
83
- if args.input_image:
84
- demo_image = Image.open(args.input_image).convert("RGB")
85
- elif args.image_url:
86
- demo_image = load_image(args.image_url)
87
- else:
88
- raise ValueError("Either --input_image or --image_url must be provided")
89
- except Exception as e:
90
- raise ValueError(f"Error loading image: {e}")
91
-
92
- poses = [controlnet_detector(demo_image)]
93
-
94
- # Generate images
95
- generators = [torch.Generator(device="cpu").manual_seed(args.seed + i) for i in range(len(poses))]
96
-
97
- output_images = generate_images(
98
- pipe,
99
- [args.prompt] * len(generators),
100
- poses,
101
- generators,
102
- [args.negative_prompt] * len(generators),
103
- args.num_steps,
104
- args.guidance_scale,
105
- args.controlnet_conditioning_scale,
106
- args.width,
107
- args.height
108
- )
109
-
110
- # Save images if save_output is True
111
- if args.save_output:
112
- os.makedirs(args.output_dir, exist_ok=True)
113
- for i, img in enumerate(tqdm(output_images, desc="Saving images")):
114
- if args.use_prompt_as_output_name:
115
- # Sanitize prompt for filename (replace spaces and special characters)
116
- sanitized_prompt = re.sub(r'[^\w\s-]', '', args.prompt).replace(' ', '_').lower()
117
- filename = f"{sanitized_prompt}_{i}.png"
118
- else:
119
- # Use UUID for filename
120
- filename = f"{uuid.uuid4()}_{i}.png"
121
- img.save(os.path.join(args.output_dir, filename))
122
-
123
- if __name__ == "__main__":
124
- parser = argparse.ArgumentParser(description="ControlNet image generation with pose detection")
125
- # Create mutually exclusive group for input_image and image_url
126
- image_group = parser.add_mutually_exclusive_group(required=True)
127
- image_group.add_argument("--input_image", type=str, default=None,
128
- help="Path to local input image (default: tests/test_data/yoga1.jpg)")
129
- image_group.add_argument("--image_url", type=str, default=None,
130
- help="URL of input image (e.g., https://huggingface.co/datasets/YiYiXu/controlnet-testing/resolve/main/yoga1.jpeg)")
131
-
132
- parser.add_argument("--config_path", type=str, default="configs/model_ckpts.yaml",
133
- help="Path to configuration YAML file")
134
- parser.add_argument("--prompt", type=str, default="a man is doing yoga",
135
- help="Text prompt for image generation")
136
- parser.add_argument("--negative_prompt", type=str,
137
- default="monochrome, lowres, bad anatomy, worst quality, low quality",
138
- help="Negative prompt for image generation")
139
- parser.add_argument("--num_steps", type=int, default=20,
140
- help="Number of inference steps")
141
- parser.add_argument("--seed", type=int, default=2,
142
- help="Random seed for generation")
143
- parser.add_argument("--width", type=int, default=512,
144
- help="Width of the generated image")
145
- parser.add_argument("--height", type=int, default=512,
146
- help="Height of the generated image")
147
- parser.add_argument("--guidance_scale", type=float, default=7.5,
148
- help="Guidance scale for prompt adherence")
149
- parser.add_argument("--controlnet_conditioning_scale", type=float, default=1.0,
150
- help="ControlNet conditioning scale")
151
- parser.add_argument("--output_dir", type=str, default="tests/test_data",
152
- help="Directory to save generated images")
153
- parser.add_argument("--use_prompt_as_output_name", action="store_true",
154
- help="Use prompt as part of output image filename")
155
- parser.add_argument("--save_output", action="store_true artr",
156
- help="Save generated images to output directory")
157
-
158
- args = parser.parse_args()
159
- infer(args)
160
-
161
- # Using image_url
162
- # python script.py \
163
- # --config_path configs/model_ckpts.yaml \
164
- # --image_url https://huggingface.co/datasets/YiYiXu/controlnet-testing/resolve/main/yoga1.jpeg \
165
- # --prompt "a man is doing yoga in a serene park" \
166
- # --negative_prompt "monochrome, lowres, bad anatomy" \
167
- # --num_steps 30 \
168
- # --seed 42 \
169
- # --width 512 \
170
- # --height 512 \
171
- # --guidance_scale 7.5 \
172
- # --controlnet_conditioning_scale 0.8 \
173
- # --output_dir "tests/test_data" \
174
- # --save_output
175
-
176
- # Using input_image
177
- # python script.py \
178
- # --config_path configs/model_ckpts.yaml \
179
- # --input_image "tests/test_data/yoga1.jpg" \
180
- # --prompt "a man is doing yoga in a serene park" \
181
- # --negative_prompt "monochrome, lowres, bad anatomy" \
182
- # --num_steps 30 \
183
- # --seed 42 \
184
- # --width 512 \
185
- # --height 512 \
186
- # --guidance_scale 7.5 \
187
- # --controlnet_conditioning_scale 0.8 \
188
- # --output_dir "tests/test_data" \
189
- # --save_output