Skip to content

C
coderai
Commit d784e7e9 authored by Your Name's avatar Your Name
Browse files
Options

Add CLI arguments for stable-diffusion-cpp-python image generation 

- Added --clip-l-path for specifying CLIP LLM model path
- Added --vae-path for specifying VAE model path
- Added --image-sample-method (default: res_multistep for Z-Image Turbo)
- Added --image-steps (default: 4 for Z-Image Turbo)
- Added --image-width (default: 512)
- Added --image-height (default: 512)
- Added --image-cfg-scale (default: 1.0 for Z-Image Turbo)
parent bae2d8cc
Hide whitespace changes
Inline Side-by-side
Showing with 101 additions and 22 deletions
coderai
View file @ d784e7e9
...@@ -4090,6 +4090,48 @@ def parse_args(): ...@@ -4090,6 +4090,48 @@ def parse_args():
default=None, default=None,
help="Model for image generation (e.g., stable-diffusion-xl-base-1.0). Can be specified multiple times for multiple models.", help="Model for image generation (e.g., stable-diffusion-xl-base-1.0). Can be specified multiple times for multiple models.",
) )
parser.add_argument(
"--clip-l-path",
type=str,
default=None,
help="Path to CLIP LLM model for image generation (stable-diffusion-cpp-python).",
)
parser.add_argument(
"--vae-path",
type=str,
default=None,
help="Path to VAE model for image generation (stable-diffusion-cpp-python).",
)
parser.add_argument(
"--image-sample-method",
type=str,
default="res_multistep",
help="Sample method for image generation (default: res_multistep for Z-Image Turbo).",
)
parser.add_argument(
"--image-steps",
type=int,
default=4,
help="Number of inference steps for image generation (default: 4 for Z-Image Turbo).",
)
parser.add_argument(
"--image-width",
type=int,
default=512,
help="Image width for generation (default: 512).",
)
parser.add_argument(
"--image-height",
type=int,
default=512,
help="Image height for generation (default: 512).",
)
parser.add_argument(
"--image-cfg-scale",
type=float,
default=1.0,
help="CFG scale for image generation (default: 1.0 for Z-Image Turbo).",
)
parser.add_argument( parser.add_argument(
"--loadall", "--loadall",
action="store_true", action="store_true",
...@@ -4611,23 +4653,38 @@ def main(): ...@@ -4611,23 +4653,38 @@ def main():
t5xxl_path = None t5xxl_path = None
vae_path = None vae_path = None
# Look for common file patterns # Use CLI arguments if provided
for f in os.listdir(model_dir) if os.path.exists(model_dir) else []: if args.clip_l_path:
if 'clip_l' in f.lower() and f.endswith(('.safetensors', '.bin')): clip_l_path = args.clip_l_path
clip_l_path = os.path.join(model_dir, f) if args.vae_path:
elif 't5xxl' in f.lower() and f.endswith(('.safetensors', '.bin')): vae_path = args.vae_path
t5xxl_path = os.path.join(model_dir, f)
elif f.endswith('.safetensors') and 'ae' in f.lower(): # Look for common file patterns only if CLI args not provided
vae_path = os.path.join(model_dir, f) if not args.clip_l_path or not args.vae_path:
for f in os.listdir(model_dir) if os.path.exists(model_dir) else []:
if not args.clip_l_path and 'clip_l' in f.lower() and f.endswith(('.safetensors', '.bin')):
clip_l_path = os.path.join(model_dir, f)
elif 't5xxl' in f.lower() and f.endswith(('.safetensors', '.bin')):
t5xxl_path = os.path.join(model_dir, f)
elif not args.vae_path and f.endswith('.safetensors') and 'ae' in f.lower():
vae_path = os.path.join(model_dir, f)
# Build kwargs based on available files # Build kwargs based on available files
sd_kwargs = {'diffusion_model_path': model_path} sd_kwargs = {'diffusion_model_path': model_path}
if clip_l_path: if args.clip_l_path:
sd_kwargs['clip_l_path'] = args.clip_l_path
elif clip_l_path:
sd_kwargs['clip_l_path'] = clip_l_path sd_kwargs['clip_l_path'] = clip_l_path
if args.vae_path:
sd_kwargs['vae_path'] = args.vae_path
elif vae_path:
sd_kwargs['vae_path'] = vae_path
if t5xxl_path: if t5xxl_path:
sd_kwargs['t5xxl_path'] = t5xxl_path sd_kwargs['t5xxl_path'] = t5xxl_path
if vae_path:
sd_kwargs['vae_path'] = vae_path # Add generation parameters from CLI args
sd_kwargs['sample_method'] = args.image_sample_method
sd_kwargs['steps'] = args.image_steps
sd_model = StableDiffusion(**sd_kwargs) sd_model = StableDiffusion(**sd_kwargs)
multi_model_manager.add_model(model_key, sd_model) multi_model_manager.add_model(model_key, sd_model)
...@@ -4935,6 +4992,13 @@ def main(): ...@@ -4935,6 +4992,13 @@ def main():
multi_model_manager.set_image_model(image_models[0], { multi_model_manager.set_image_model(image_models[0], {
'ctx': args.vision_ctx, 'ctx': args.vision_ctx,
'offload': args.vision_offload, 'offload': args.vision_offload,
'clip_l_path': args.clip_l_path,
'vae_path': args.vae_path,
'sample_method': args.image_sample_method,
'steps': args.image_steps,
'width': args.image_width,
'height': args.image_height,
'cfg_scale': args.image_cfg_scale,
}) })
# Register all image models # Register all image models
for img_m in image_models[1:]: for img_m in image_models[1:]:
...@@ -5049,23 +5113,38 @@ def main(): ...@@ -5049,23 +5113,38 @@ def main():
t5xxl_path = None t5xxl_path = None
vae_path = None vae_path = None
# Look for common file patterns # Use CLI arguments if provided
for f in os.listdir(model_dir) if os.path.exists(model_dir) else []: if args.clip_l_path:
if 'clip_l' in f.lower() and f.endswith(('.safetensors', '.bin')): clip_l_path = args.clip_l_path
clip_l_path = os.path.join(model_dir, f) if args.vae_path:
elif 't5xxl' in f.lower() and f.endswith(('.safetensors', '.bin')): vae_path = args.vae_path
t5xxl_path = os.path.join(model_dir, f)
elif f.endswith('.safetensors') and 'ae' in f.lower(): # Look for common file patterns only if CLI args not provided
vae_path = os.path.join(model_dir, f) if not args.clip_l_path or not args.vae_path:
for f in os.listdir(model_dir) if os.path.exists(model_dir) else []:
if not args.clip_l_path and 'clip_l' in f.lower() and f.endswith(('.safetensors', '.bin')):
clip_l_path = os.path.join(model_dir, f)
elif 't5xxl' in f.lower() and f.endswith(('.safetensors', '.bin')):
t5xxl_path = os.path.join(model_dir, f)
elif not args.vae_path and f.endswith('.safetensors') and 'ae' in f.lower():
vae_path = os.path.join(model_dir, f)
# Build kwargs based on available files # Build kwargs based on available files
sd_kwargs = {'diffusion_model_path': model_path} sd_kwargs = {'diffusion_model_path': model_path}
if clip_l_path: if args.clip_l_path:
sd_kwargs['clip_l_path'] = args.clip_l_path
elif clip_l_path:
sd_kwargs['clip_l_path'] = clip_l_path sd_kwargs['clip_l_path'] = clip_l_path
if args.vae_path:
sd_kwargs['vae_path'] = args.vae_path
elif vae_path:
sd_kwargs['vae_path'] = vae_path
if t5xxl_path: if t5xxl_path:
sd_kwargs['t5xxl_path'] = t5xxl_path sd_kwargs['t5xxl_path'] = t5xxl_path
if vae_path:
sd_kwargs['vae_path'] = vae_path # Add generation parameters from CLI args
sd_kwargs['sample_method'] = args.image_sample_method
sd_kwargs['steps'] = args.image_steps
sd_model = StableDiffusion(**sd_kwargs) sd_model = StableDiffusion(**sd_kwargs)
multi_model_manager.add_model(model_key, sd_model) multi_model_manager.add_model(model_key, sd_model)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment