Spaces:

HTBris
/

hyp4thi4

Running

hyp4thi4 / generate.py

Check for unet safetensors before using local hf-model, otherwise load from HF Hub

a5719d7 verified 3 days ago

16.5 kB

#!/usr/bin/env python3
"""
Standalone Offline CPU Image Generator for ModelsLab/indigo-furry-mix-v65.
Supports:
- 100% offline local loading from ./hf-model (and ./hf-model/lcm-lora for --lcm)
- Standard schedulers (Euler, Euler a, DDIM, DPM++ 2M, DPM++ 2M Karras, PNDM, Heun, LMS, LCM)
- Optional LCM-LoRA ultra-fast 4-8 step generation on CPU (--lcm)
- Single-run CLI mode (--prompt "...") or interactive loop mode (--interactive or omit --prompt)
"""

import argparse
import random
import shlex
from datetime import datetime
from pathlib import Path

import torch
from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    HeunDiscreteScheduler,
    LCMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPipeline,
)

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_PATH = SCRIPT_DIR / "hf-model"
DEFAULT_LCM_LORA_PATH = DEFAULT_MODEL_PATH / "lcm-lora"
HF_BASE_MODEL_ID = "ModelsLab/indigo-furry-mix-v65"
HF_LCM_LORA_ID = "latent-consistency/lcm-lora-sdv1-5"

DEFAULT_NEGATIVE_PROMPT = (
    "lowres, bad anatomy, bad hands, text, error, missing fingers, "
    "extra digit, fewer digits, cropped, worst quality, low quality, "
    "normal quality, jpeg artifacts, signature, watermark, username, blurry"
)


def configure_scheduler(pipe: StableDiffusionPipeline, scheduler_name: str) -> None:
    """Replaces the pipeline's scheduler based on the user's CLI choice."""
    name = scheduler_name.strip().lower().replace("-", "_").replace(" ", "_")
    config = pipe.scheduler.config

    if name in ("ddim",):
        pipe.scheduler = DDIMScheduler.from_config(config)
    elif name in ("euler",):
        pipe.scheduler = EulerDiscreteScheduler.from_config(config)
    elif name in ("euler_a", "euler_ancestral"):
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(config)
    elif name in ("dpm++", "dpm++_2m", "dpmpp_2m"):
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(config)
    elif name in ("dpm++_karras", "dpm++_2m_karras", "dpmpp_2m_karras", "karras"):
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            config, use_karras_sigmas=True
        )
    elif name in ("dpm++_sde", "dpmpp_sde"):
        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(config)
    elif name in ("dpm++_sde_karras", "dpmpp_sde_karras"):
        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
            config, use_karras_sigmas=True
        )
    elif name in ("pndm",):
        pipe.scheduler = PNDMScheduler.from_config(config)
    elif name in ("heun",):
        pipe.scheduler = HeunDiscreteScheduler.from_config(config)
    elif name in ("lms",):
        pipe.scheduler = LMSDiscreteScheduler.from_config(config)
    elif name in ("lcm",):
        pipe.scheduler = LCMScheduler.from_config(config)
    else:
        supported = (
            "ddim, euler, euler_a, dpm++, dpm++_karras, "
            "dpm++_sde, dpm++_sde_karras, pndm, heun, lms, lcm"
        )
        raise ValueError(
            f"Unknown scheduler '{scheduler_name}'. Supported values: {supported}"
        )


def resolve_seed(seed_arg: str) -> int:
    """Parses '--seed r' (or 'random' / '-1') into a random 32-bit integer, or converts numeric string."""
    val = str(seed_arg).strip().lower()
    if val in ("r", "random", "-1", "none", ""):
        return random.randint(0, 2**32 - 1)
    try:
        return int(val)
    except ValueError as err:
        raise ValueError(
            f"Invalid seed '{seed_arg}'. Use an integer (e.g. 12345) or 'r' for random."
        ) from err


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate images offline on CPU using local hf-model",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        default=None,
        help="Positive text prompt. If omitted, starts in interactive mode.",
    )
    parser.add_argument(
        "-i",
        "--interactive",
        action="store_true",
        help="Keep the model loaded in memory and enter an interactive prompt loop.",
    )
    parser.add_argument(
        "-n",
        "--negative-prompt",
        type=str,
        default=DEFAULT_NEGATIVE_PROMPT,
        help="Negative prompt (things to avoid in the image).",
    )
    parser.add_argument(
        "--steps",
        type=int,
        default=None,
        help="Number of inference steps (defaults to 4 when --lcm is used, otherwise 20).",
    )
    parser.add_argument(
        "--cfg",
        type=float,
        default=None,
        help="Guidance scale / CFG (defaults to 1.5 when --lcm is used, otherwise 7.0).",
    )
    parser.add_argument(
        "--seed",
        type=str,
        default="r",
        help="Seed number for reproducibility, or 'r' for a random seed.",
    )
    parser.add_argument(
        "--scheduler",
        type=str,
        default=None,
        help=(
            "Sampler/scheduler to use: ddim, euler, euler_a, dpm++, dpm++_karras, "
            "dpm++_sde, pndm, heun, lms, lcm (defaults to 'lcm' with --lcm, else 'dpm++_karras')."
        ),
    )
    parser.add_argument(
        "--lcm",
        action="store_true",
        help="Enable LCM-LoRA mode for ultra-fast 4-8 step generation on CPU.",
    )
    parser.add_argument(
        "--model-path",
        type=str,
        default=str(DEFAULT_MODEL_PATH),
        help="Path to the local model directory.",
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        default=str(DEFAULT_LCM_LORA_PATH),
        help="Path to the local LCM-LoRA directory.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=512,
        help="Image width in pixels (must be a multiple of 8).",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=512,
        help="Image height in pixels (must be a multiple of 8).",
    )
    parser.add_argument(
        "--count",
        "--num-images",
        "-c",
        dest="num_images",
        type=int,
        default=1,
        help="Number of images to generate per prompt.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default=None,
        help="Custom output file path. If omitted, saves to outputs/<timestamp>_seed<seed>.png.",
    )
    parser.add_argument(
        "--threads",
        type=int,
        default=None,
        help="Optional number of CPU threads for PyTorch to use.",
    )
    return parser.parse_args()


def generate_one(
    pipe: StableDiffusionPipeline,
    prompt: str,
    negative_prompt: str,
    steps: int,
    cfg: float,
    seed_arg: str,
    scheduler_name: str,
    lcm_active: bool,
    width: int,
    height: int,
    num_images: int = 1,
    output_override: str | None = None,
) -> list[Path]:
    """Runs image generation (1 or more images) and saves the results to disk."""
    configure_scheduler(pipe, scheduler_name)
    base_seed = resolve_seed(seed_arg)
    seeds = [base_seed + i for i in range(num_images)]

    print("\n--- Generation Config ---")
    print(f"Prompt          : {prompt}")
    print(f"Negative Prompt : {negative_prompt}")
    print(f"Scheduler       : {scheduler_name} ({pipe.scheduler.__class__.__name__})")
    print(f"LCM-LoRA Active : {lcm_active}")
    print(f"Steps           : {steps}")
    print(f"CFG Scale       : {cfg}")
    print(f"Images (Count)  : {num_images}")
    print(f"Seed(s)         : {seeds if num_images > 1 else base_seed}")
    print(f"Resolution      : {width}x{height}")
    print("-------------------------\n")

    generators = [
        torch.Generator(device="cpu").manual_seed(s) for s in seeds
    ]

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        guidance_scale=cfg,
        width=width,
        height=height,
        num_images_per_prompt=num_images,
        generator=generators,
    )

    saved_paths: list[Path] = []
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    for idx, (image, s) in enumerate(zip(result.images, seeds), start=1):
        if output_override:
            orig = Path(output_override)
            orig.parent.mkdir(parents=True, exist_ok=True)
            if num_images == 1:
                output_path = orig
            else:
                output_path = orig.with_name(f"{orig.stem}_{idx}{orig.suffix}")
        else:
            output_dir = SCRIPT_DIR / "outputs"
            output_dir.mkdir(parents=True, exist_ok=True)
            if num_images == 1:
                output_path = output_dir / f"{timestamp}_seed{s}.png"
            else:
                output_path = output_dir / f"{timestamp}_{idx}_seed{s}.png"

        image.save(output_path)
        saved_paths.append(output_path)
        print(f"[+] Saved image ({idx}/{num_images}) to: {output_path.resolve()}")

    return saved_paths


def parse_inline_overrides(
    line: str,
    default_steps: int,
    default_cfg: float,
    default_seed: str,
    default_scheduler: str,
    default_width: int,
    default_height: int,
    default_num_images: int,
) -> tuple[str, int, float, str, str, int, int, int]:
    """
    Allows interactive prompts to optionally include inline CLI flags like:
    `a cool wolf --steps 6 --cfg 1.8 --seed 42 --scheduler euler --count 2`
    """
    if " --" not in line and " -c " not in line:
        return (
            line.strip(),
            default_steps,
            default_cfg,
            default_seed,
            default_scheduler,
            default_width,
            default_height,
            default_num_images,
        )

    tokens = shlex.split(line)
    prompt_parts = []
    steps = default_steps
    cfg = default_cfg
    seed = default_seed
    scheduler = default_scheduler
    width = default_width
    height = default_height
    num_images = default_num_images

    i = 0
    while i < len(tokens):
        tok = tokens[i]
        if tok == "--steps" and i + 1 < len(tokens):
            steps = int(tokens[i + 1])
            i += 2
        elif tok == "--cfg" and i + 1 < len(tokens):
            cfg = float(tokens[i + 1])
            i += 2
        elif tok == "--seed" and i + 1 < len(tokens):
            seed = tokens[i + 1]
            i += 2
        elif tok == "--scheduler" and i + 1 < len(tokens):
            scheduler = tokens[i + 1]
            i += 2
        elif tok == "--width" and i + 1 < len(tokens):
            width = int(tokens[i + 1])
            i += 2
        elif tok == "--height" and i + 1 < len(tokens):
            height = int(tokens[i + 1])
            i += 2
        elif tok in ("--count", "--num-images", "-c") and i + 1 < len(tokens):
            num_images = int(tokens[i + 1])
            i += 2
        else:
            prompt_parts.append(tok)
            i += 1

    return (
        " ".join(prompt_parts).strip(),
        steps,
        cfg,
        seed,
        scheduler,
        width,
        height,
        num_images,
    )


def main() -> None:
    args = parse_args()

    if args.threads is not None and args.threads > 0:
        torch.set_num_threads(args.threads)

    steps = args.steps if args.steps is not None else (4 if args.lcm else 20)
    cfg = args.cfg if args.cfg is not None else (1.5 if args.lcm else 7.0)
    scheduler_name = (
        args.scheduler
        if args.scheduler is not None
        else ("lcm" if args.lcm else "dpm++_karras")
    )

    model_path = Path(args.model_path).resolve()
    unet_file = model_path / "unet" / "diffusion_pytorch_model.fp16.safetensors"
    if model_path.exists() and unet_file.exists():
        model_source = str(model_path)
        local_model_only = True
        print(f"[*] Loading local model from '{model_path}' onto CPU (offline mode)...")
    else:
        model_source = HF_BASE_MODEL_ID
        local_model_only = False
        print(f"[*] Local model weights not found at '{model_path}'. Loading '{model_source}' from Hugging Face Hub...")

    # Load using variant="fp16" (matching *.fp16.safetensors)
    # and upcast in-memory to torch.float32 for CPU compatibility.
    pipe = StableDiffusionPipeline.from_pretrained(
        model_source,
        torch_dtype=torch.float32,
        variant="fp16",
        use_safetensors=True,
        safety_checker=None,
        requires_safety_checker=False,
        local_files_only=local_model_only,
    )

    if args.lcm:
        lora_path = Path(args.lora_path).resolve()
        lora_file = lora_path / "pytorch_lora_weights.safetensors"
        if lora_file.exists():
            lora_source = str(lora_path)
            local_lora_only = True
            print(f"[*] Loading and fusing local LCM-LoRA from '{lora_path}'...")
        else:
            lora_source = HF_LCM_LORA_ID
            local_lora_only = False
            print(f"[*] Local LCM-LoRA not found at '{lora_file}'. Loading '{lora_source}' from Hugging Face Hub...")
        pipe.load_lora_weights(
            lora_source,
            weight_name="pytorch_lora_weights.safetensors",
            local_files_only=local_lora_only,
        )
        pipe.fuse_lora()

    # CPU memory & speed optimizations
    pipe.to("cpu")
    pipe.enable_attention_slicing()
    if hasattr(pipe, "enable_vae_slicing"):
        pipe.enable_vae_slicing()

    # Run initial prompt if provided on command line
    if args.prompt:
        generate_one(
            pipe=pipe,
            prompt=args.prompt,
            negative_prompt=args.negative_prompt,
            steps=steps,
            cfg=cfg,
            seed_arg=args.seed,
            scheduler_name=scheduler_name,
            lcm_active=args.lcm,
            width=args.width,
            height=args.height,
            num_images=args.num_images,
            output_override=args.output,
        )

    # Enter interactive loop if --interactive was passed OR no --prompt was given
    if args.interactive or not args.prompt:
        print("\n========================================================")
        print(" Interactive Mode Ready! (Model stays loaded in RAM)")
        print(" Type a prompt and press Enter to generate.")
        print(" Tip: You can also append flags inline, e.g.:")
        print("   cyberpunk wolf --steps 4 --cfg 1.5 --seed 42 --count 2")
        print(" Type 'q' or 'exit' to quit.")
        print("========================================================\n")

        while True:
            try:
                user_input = input("Prompt > ").strip()
            except (EOFError, KeyboardInterrupt):
                print("\nExiting interactive mode.")
                break

            if not user_input:
                continue
            if user_input.lower() in ("q", "quit", "exit"):
                print("Exiting interactive mode.")
                break

            try:
                (
                    p_text,
                    p_steps,
                    p_cfg,
                    p_seed,
                    p_sched,
                    p_width,
                    p_height,
                    p_count,
                ) = parse_inline_overrides(
                    user_input,
                    default_steps=steps,
                    default_cfg=cfg,
                    default_seed=args.seed,
                    default_scheduler=scheduler_name,
                    default_width=args.width,
                    default_height=args.height,
                    default_num_images=args.num_images,
                )
                if not p_text:
                    print("[!] Please provide a non-empty prompt.")
                    continue

                generate_one(
                    pipe=pipe,
                    prompt=p_text,
                    negative_prompt=args.negative_prompt,
                    steps=p_steps,
                    cfg=p_cfg,
                    seed_arg=p_seed,
                    scheduler_name=p_sched,
                    lcm_active=args.lcm,
                    width=p_width,
                    height=p_height,
                    num_images=p_count,
                    output_override=None,
                )
            except Exception as err:
                print(f"[!] Error during generation: {err}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	Standalone Offline CPU Image Generator for ModelsLab/indigo-furry-mix-v65.

	Supports:
	- 100% offline local loading from ./hf-model (and ./hf-model/lcm-lora for --lcm)
	- Standard schedulers (Euler, Euler a, DDIM, DPM++ 2M, DPM++ 2M Karras, PNDM, Heun, LMS, LCM)
	- Optional LCM-LoRA ultra-fast 4-8 step generation on CPU (--lcm)
	- Single-run CLI mode (--prompt "...") or interactive loop mode (--interactive or omit --prompt)
	"""

	import argparse
	import random
	import shlex
	from datetime import datetime
	from pathlib import Path

	import torch
	from diffusers import (
	DDIMScheduler,
	DPMSolverMultistepScheduler,
	DPMSolverSinglestepScheduler,
	EulerAncestralDiscreteScheduler,
	EulerDiscreteScheduler,
	HeunDiscreteScheduler,
	LCMScheduler,
	LMSDiscreteScheduler,
	PNDMScheduler,
	StableDiffusionPipeline,
	)

	SCRIPT_DIR = Path(__file__).resolve().parent
	DEFAULT_MODEL_PATH = SCRIPT_DIR / "hf-model"
	DEFAULT_LCM_LORA_PATH = DEFAULT_MODEL_PATH / "lcm-lora"
	HF_BASE_MODEL_ID = "ModelsLab/indigo-furry-mix-v65"
	HF_LCM_LORA_ID = "latent-consistency/lcm-lora-sdv1-5"

	DEFAULT_NEGATIVE_PROMPT = (
	"lowres, bad anatomy, bad hands, text, error, missing fingers, "
	"extra digit, fewer digits, cropped, worst quality, low quality, "
	"normal quality, jpeg artifacts, signature, watermark, username, blurry"
	)


	def configure_scheduler(pipe: StableDiffusionPipeline, scheduler_name: str) -> None:
	"""Replaces the pipeline's scheduler based on the user's CLI choice."""
	name = scheduler_name.strip().lower().replace("-", "_").replace(" ", "_")
	config = pipe.scheduler.config

	if name in ("ddim",):
	pipe.scheduler = DDIMScheduler.from_config(config)
	elif name in ("euler",):
	pipe.scheduler = EulerDiscreteScheduler.from_config(config)
	elif name in ("euler_a", "euler_ancestral"):
	pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(config)
	elif name in ("dpm++", "dpm++_2m", "dpmpp_2m"):
	pipe.scheduler = DPMSolverMultistepScheduler.from_config(config)
	elif name in ("dpm++_karras", "dpm++_2m_karras", "dpmpp_2m_karras", "karras"):
	pipe.scheduler = DPMSolverMultistepScheduler.from_config(
	config, use_karras_sigmas=True
	)
	elif name in ("dpm++_sde", "dpmpp_sde"):
	pipe.scheduler = DPMSolverSinglestepScheduler.from_config(config)
	elif name in ("dpm++_sde_karras", "dpmpp_sde_karras"):
	pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
	config, use_karras_sigmas=True
	)
	elif name in ("pndm",):
	pipe.scheduler = PNDMScheduler.from_config(config)
	elif name in ("heun",):
	pipe.scheduler = HeunDiscreteScheduler.from_config(config)
	elif name in ("lms",):
	pipe.scheduler = LMSDiscreteScheduler.from_config(config)
	elif name in ("lcm",):
	pipe.scheduler = LCMScheduler.from_config(config)
	else:
	supported = (
	"ddim, euler, euler_a, dpm++, dpm++_karras, "
	"dpm++_sde, dpm++_sde_karras, pndm, heun, lms, lcm"
	)
	raise ValueError(
	f"Unknown scheduler '{scheduler_name}'. Supported values: {supported}"
	)


	def resolve_seed(seed_arg: str) -> int:
	"""Parses '--seed r' (or 'random' / '-1') into a random 32-bit integer, or converts numeric string."""
	val = str(seed_arg).strip().lower()
	if val in ("r", "random", "-1", "none", ""):
	return random.randint(0, 2**32 - 1)
	try:
	return int(val)
	except ValueError as err:
	raise ValueError(
	f"Invalid seed '{seed_arg}'. Use an integer (e.g. 12345) or 'r' for random."
	) from err


	def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(
	description="Generate images offline on CPU using local hf-model",
	formatter_class=argparse.ArgumentDefaultsHelpFormatter,
	)
	parser.add_argument(
	"-p",
	"--prompt",
	type=str,
	default=None,
	help="Positive text prompt. If omitted, starts in interactive mode.",
	)
	parser.add_argument(
	"-i",
	"--interactive",
	action="store_true",
	help="Keep the model loaded in memory and enter an interactive prompt loop.",
	)
	parser.add_argument(
	"-n",
	"--negative-prompt",
	type=str,
	default=DEFAULT_NEGATIVE_PROMPT,
	help="Negative prompt (things to avoid in the image).",
	)
	parser.add_argument(
	"--steps",
	type=int,
	default=None,
	help="Number of inference steps (defaults to 4 when --lcm is used, otherwise 20).",
	)
	parser.add_argument(
	"--cfg",
	type=float,
	default=None,
	help="Guidance scale / CFG (defaults to 1.5 when --lcm is used, otherwise 7.0).",
	)
	parser.add_argument(
	"--seed",
	type=str,
	default="r",
	help="Seed number for reproducibility, or 'r' for a random seed.",
	)
	parser.add_argument(
	"--scheduler",
	type=str,
	default=None,
	help=(
	"Sampler/scheduler to use: ddim, euler, euler_a, dpm++, dpm++_karras, "
	"dpm++_sde, pndm, heun, lms, lcm (defaults to 'lcm' with --lcm, else 'dpm++_karras')."
	),
	)
	parser.add_argument(
	"--lcm",
	action="store_true",
	help="Enable LCM-LoRA mode for ultra-fast 4-8 step generation on CPU.",
	)
	parser.add_argument(
	"--model-path",
	type=str,
	default=str(DEFAULT_MODEL_PATH),
	help="Path to the local model directory.",
	)
	parser.add_argument(
	"--lora-path",
	type=str,
	default=str(DEFAULT_LCM_LORA_PATH),
	help="Path to the local LCM-LoRA directory.",
	)
	parser.add_argument(
	"--width",
	type=int,
	default=512,
	help="Image width in pixels (must be a multiple of 8).",
	)
	parser.add_argument(
	"--height",
	type=int,
	default=512,
	help="Image height in pixels (must be a multiple of 8).",
	)
	parser.add_argument(
	"--count",
	"--num-images",
	"-c",
	dest="num_images",
	type=int,
	default=1,
	help="Number of images to generate per prompt.",
	)
	parser.add_argument(
	"-o",
	"--output",
	type=str,
	default=None,
	help="Custom output file path. If omitted, saves to outputs/<timestamp>_seed<seed>.png.",
	)
	parser.add_argument(
	"--threads",
	type=int,
	default=None,
	help="Optional number of CPU threads for PyTorch to use.",
	)
	return parser.parse_args()


	def generate_one(
	pipe: StableDiffusionPipeline,
	prompt: str,
	negative_prompt: str,
	steps: int,
	cfg: float,
	seed_arg: str,
	scheduler_name: str,
	lcm_active: bool,
	width: int,
	height: int,
	num_images: int = 1,
	output_override: str \| None = None,
	) -> list[Path]:
	"""Runs image generation (1 or more images) and saves the results to disk."""
	configure_scheduler(pipe, scheduler_name)
	base_seed = resolve_seed(seed_arg)
	seeds = [base_seed + i for i in range(num_images)]

	print("\n--- Generation Config ---")
	print(f"Prompt : {prompt}")
	print(f"Negative Prompt : {negative_prompt}")
	print(f"Scheduler : {scheduler_name} ({pipe.scheduler.__class__.__name__})")
	print(f"LCM-LoRA Active : {lcm_active}")
	print(f"Steps : {steps}")
	print(f"CFG Scale : {cfg}")
	print(f"Images (Count) : {num_images}")
	print(f"Seed(s) : {seeds if num_images > 1 else base_seed}")
	print(f"Resolution : {width}x{height}")
	print("-------------------------\n")

	generators = [
	torch.Generator(device="cpu").manual_seed(s) for s in seeds
	]

	result = pipe(
	prompt=prompt,
	negative_prompt=negative_prompt,
	num_inference_steps=steps,
	guidance_scale=cfg,
	width=width,
	height=height,
	num_images_per_prompt=num_images,
	generator=generators,
	)

	saved_paths: list[Path] = []
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

	for idx, (image, s) in enumerate(zip(result.images, seeds), start=1):
	if output_override:
	orig = Path(output_override)
	orig.parent.mkdir(parents=True, exist_ok=True)
	if num_images == 1:
	output_path = orig
	else:
	output_path = orig.with_name(f"{orig.stem}_{idx}{orig.suffix}")
	else:
	output_dir = SCRIPT_DIR / "outputs"
	output_dir.mkdir(parents=True, exist_ok=True)
	if num_images == 1:
	output_path = output_dir / f"{timestamp}_seed{s}.png"
	else:
	output_path = output_dir / f"{timestamp}_{idx}_seed{s}.png"

	image.save(output_path)
	saved_paths.append(output_path)
	print(f"[+] Saved image ({idx}/{num_images}) to: {output_path.resolve()}")

	return saved_paths


	def parse_inline_overrides(
	line: str,
	default_steps: int,
	default_cfg: float,
	default_seed: str,
	default_scheduler: str,
	default_width: int,
	default_height: int,
	default_num_images: int,
	) -> tuple[str, int, float, str, str, int, int, int]:
	"""
	Allows interactive prompts to optionally include inline CLI flags like:
	`a cool wolf --steps 6 --cfg 1.8 --seed 42 --scheduler euler --count 2`
	"""
	if " --" not in line and " -c " not in line:
	return (
	line.strip(),
	default_steps,
	default_cfg,
	default_seed,
	default_scheduler,
	default_width,
	default_height,
	default_num_images,
	)

	tokens = shlex.split(line)
	prompt_parts = []
	steps = default_steps
	cfg = default_cfg
	seed = default_seed
	scheduler = default_scheduler
	width = default_width
	height = default_height
	num_images = default_num_images

	i = 0
	while i < len(tokens):
	tok = tokens[i]
	if tok == "--steps" and i + 1 < len(tokens):
	steps = int(tokens[i + 1])
	i += 2
	elif tok == "--cfg" and i + 1 < len(tokens):
	cfg = float(tokens[i + 1])
	i += 2
	elif tok == "--seed" and i + 1 < len(tokens):
	seed = tokens[i + 1]
	i += 2
	elif tok == "--scheduler" and i + 1 < len(tokens):
	scheduler = tokens[i + 1]
	i += 2
	elif tok == "--width" and i + 1 < len(tokens):
	width = int(tokens[i + 1])
	i += 2
	elif tok == "--height" and i + 1 < len(tokens):
	height = int(tokens[i + 1])
	i += 2
	elif tok in ("--count", "--num-images", "-c") and i + 1 < len(tokens):
	num_images = int(tokens[i + 1])
	i += 2
	else:
	prompt_parts.append(tok)
	i += 1

	return (
	" ".join(prompt_parts).strip(),
	steps,
	cfg,
	seed,
	scheduler,
	width,
	height,
	num_images,
	)


	def main() -> None:
	args = parse_args()

	if args.threads is not None and args.threads > 0:
	torch.set_num_threads(args.threads)

	steps = args.steps if args.steps is not None else (4 if args.lcm else 20)
	cfg = args.cfg if args.cfg is not None else (1.5 if args.lcm else 7.0)
	scheduler_name = (
	args.scheduler
	if args.scheduler is not None
	else ("lcm" if args.lcm else "dpm++_karras")
	)

	model_path = Path(args.model_path).resolve()
	unet_file = model_path / "unet" / "diffusion_pytorch_model.fp16.safetensors"
	if model_path.exists() and unet_file.exists():
	model_source = str(model_path)
	local_model_only = True
	print(f"[*] Loading local model from '{model_path}' onto CPU (offline mode)...")
	else:
	model_source = HF_BASE_MODEL_ID
	local_model_only = False
	print(f"[*] Local model weights not found at '{model_path}'. Loading '{model_source}' from Hugging Face Hub...")

	# Load using variant="fp16" (matching *.fp16.safetensors)
	# and upcast in-memory to torch.float32 for CPU compatibility.
	pipe = StableDiffusionPipeline.from_pretrained(
	model_source,
	torch_dtype=torch.float32,
	variant="fp16",
	use_safetensors=True,
	safety_checker=None,
	requires_safety_checker=False,
	local_files_only=local_model_only,
	)

	if args.lcm:
	lora_path = Path(args.lora_path).resolve()
	lora_file = lora_path / "pytorch_lora_weights.safetensors"
	if lora_file.exists():
	lora_source = str(lora_path)
	local_lora_only = True
	print(f"[*] Loading and fusing local LCM-LoRA from '{lora_path}'...")
	else:
	lora_source = HF_LCM_LORA_ID
	local_lora_only = False
	print(f"[*] Local LCM-LoRA not found at '{lora_file}'. Loading '{lora_source}' from Hugging Face Hub...")
	pipe.load_lora_weights(
	lora_source,
	weight_name="pytorch_lora_weights.safetensors",
	local_files_only=local_lora_only,
	)
	pipe.fuse_lora()

	# CPU memory & speed optimizations
	pipe.to("cpu")
	pipe.enable_attention_slicing()
	if hasattr(pipe, "enable_vae_slicing"):
	pipe.enable_vae_slicing()

	# Run initial prompt if provided on command line
	if args.prompt:
	generate_one(
	pipe=pipe,
	prompt=args.prompt,
	negative_prompt=args.negative_prompt,
	steps=steps,
	cfg=cfg,
	seed_arg=args.seed,
	scheduler_name=scheduler_name,
	lcm_active=args.lcm,
	width=args.width,
	height=args.height,
	num_images=args.num_images,
	output_override=args.output,
	)

	# Enter interactive loop if --interactive was passed OR no --prompt was given
	if args.interactive or not args.prompt:
	print("\n========================================================")
	print(" Interactive Mode Ready! (Model stays loaded in RAM)")
	print(" Type a prompt and press Enter to generate.")
	print(" Tip: You can also append flags inline, e.g.:")
	print(" cyberpunk wolf --steps 4 --cfg 1.5 --seed 42 --count 2")
	print(" Type 'q' or 'exit' to quit.")
	print("========================================================\n")

	while True:
	try:
	user_input = input("Prompt > ").strip()
	except (EOFError, KeyboardInterrupt):
	print("\nExiting interactive mode.")
	break

	if not user_input:
	continue
	if user_input.lower() in ("q", "quit", "exit"):
	print("Exiting interactive mode.")
	break

	try:
	(
	p_text,
	p_steps,
	p_cfg,
	p_seed,
	p_sched,
	p_width,
	p_height,
	p_count,
	) = parse_inline_overrides(
	user_input,
	default_steps=steps,
	default_cfg=cfg,
	default_seed=args.seed,
	default_scheduler=scheduler_name,
	default_width=args.width,
	default_height=args.height,
	default_num_images=args.num_images,
	)
	if not p_text:
	print("[!] Please provide a non-empty prompt.")
	continue

	generate_one(
	pipe=pipe,
	prompt=p_text,
	negative_prompt=args.negative_prompt,
	steps=p_steps,
	cfg=p_cfg,
	seed_arg=p_seed,
	scheduler_name=p_sched,
	lcm_active=args.lcm,
	width=p_width,
	height=p_height,
	num_images=p_count,
	output_override=None,
	)
	except Exception as err:
	print(f"[!] Error during generation: {err}")


	if __name__ == "__main__":
	main()