| """ |
| SDXS-1B image generation from prompts.txt. |
| |
| Usage: |
| python generate.py # generates grid from prompts.txt |
| python generate.py --prompts my_list.txt # custom prompts file |
| python generate.py --prompt "your text" # single prompt (без файла) |
| python generate.py --index 0 # только промпт #0 из prompts.txt |
| """ |
| import argparse |
| import os |
| import math |
| import textwrap |
|
|
| import torch |
| import numpy as np |
|
|
| from pipeline_sdxs import SdxsPipeline |
|
|
| NEGATIVE_PROMPT = ( |
| "bad quality grainy image with low details, incomplete text, " |
| "despite numerous technical flaws and distorted figures" |
| ) |
|
|
| DEFAULT_PROMPTS_FILE = "prompts.txt" |
| GRID_OUTPUT = "media/result_grid.jpg" |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="SDXS-1B generation") |
| parser.add_argument("--prompts", type=str, default=None, help="prompts file") |
| parser.add_argument("--prompt", type=str, default=None, help="single prompt") |
| parser.add_argument("--index", type=int, default=None, help="generate only prompt #N from file") |
| parser.add_argument("--steps", type=int, default=40) |
| parser.add_argument("--guidance", type=float, default=5.0) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--width", type=int, default=1024) |
| parser.add_argument("--height", type=int, default=1408) |
| parser.add_argument("--cols", type=int, default=4, help="колонки в гриде") |
| parser.add_argument("--output", type=str, default=None, |
| help="папка для отдельных картинок (по умолчанию не сохраняются)") |
| args = parser.parse_args() |
|
|
| |
| prompts = None |
| if args.prompt: |
| prompts = [args.prompt] |
| else: |
| path = args.prompts or DEFAULT_PROMPTS_FILE |
| if not os.path.exists(path): |
| print(f"Файл промптов не найден: {path}") |
| return |
| with open(path, encoding="utf-8") as f: |
| prompts = [l.strip() for l in f if l.strip()] |
|
|
| if args.index is not None: |
| if args.index >= len(prompts): |
| print(f"Нет промпта #{args.index} (всего {len(prompts)})") |
| return |
| prompts = [prompts[args.index]] |
|
|
| if not prompts: |
| print("Нет промптов") |
| return |
|
|
| print(f"Промптов: {len(prompts)}") |
|
|
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| pipe = SdxsPipeline.from_pretrained( |
| os.path.dirname(os.path.abspath(__file__)), |
| torch_dtype=torch.float16, |
| ).to(device) |
|
|
| |
| images = [] |
| for i, prompt in enumerate(prompts): |
| print(f"[{i+1}/{len(prompts)}] {prompt[:80]}...") |
| image = pipe( |
| prompt=prompt, |
| negative_prompt=NEGATIVE_PROMPT, |
| guidance_scale=args.guidance, |
| width=args.width, |
| height=args.height, |
| seed=args.seed, |
| num_inference_steps=args.steps, |
| )[0][0] |
| images.append(image) |
|
|
| if args.output: |
| os.makedirs(args.output, exist_ok=True) |
| path = os.path.join(args.output, f"{i+1:03d}.png") |
| image.save(path) |
| print(f" saved {path}") |
|
|
| |
| os.makedirs(os.path.dirname(GRID_OUTPUT), exist_ok=True) |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| cols = min(args.cols, len(images)) |
| rows = math.ceil(len(images) / cols) |
| fig, axes = plt.subplots(rows, cols, figsize=(cols * 4, rows * 4.5), constrained_layout=True) |
| axes = list(np.array(axes).flatten())[:len(images)] |
|
|
| for i, (img, prompt) in enumerate(zip(images, prompts)): |
| ax = axes[i] |
| ax.imshow(img) |
| ax.axis("off") |
| ax.set_aspect("equal") |
| text = (prompt[:200] + "…") if len(prompt) > 200 else prompt |
| lines = textwrap.wrap(text, width=35) |
| while len(lines) < 4: |
| lines.append("") |
| ax.set_title("\n".join(lines), fontsize=9, pad=8) |
|
|
| plt.savefig(GRID_OUTPUT, bbox_inches="tight", dpi=150, format="jpeg") |
| print(f"grid -> {GRID_OUTPUT}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
