| """ |
| SDXS-2B image generation from prompts.txt. |
| Usage: |
| python generate.py # generates from prompts.txt |
| python generate.py --prompts my_list.txt # custom file |
| python generate.py --prompt "your text" # single prompt |
| """ |
| import argparse |
| import os |
| import sys |
| import math |
| import textwrap |
|
|
| import torch |
| import numpy as np |
|
|
| from pipeline_sdxs import SdxsPipeline |
|
|
| NEGATIVE_PROMPT = ( |
| "low quality, bad quality, blurry, sketch, sepia, text, " |
| "bad anatomy, bad proportions, bad hands, missing fingers, child drawing" |
| ) |
|
|
| DEFAULT_PROMPTS_FILE = "prompts.txt" |
| MEDIA_DIR = "media" |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="SDXS-2B 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("--steps", type=int, default=40) |
| parser.add_argument("--guidance", type=float, default=4.0) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--width", type=int, default=832) |
| parser.add_argument("--height", type=int, default=1152) |
| parser.add_argument("--output", type=str, default=None, |
| help="output dir for individual images (optional)") |
| args = parser.parse_args() |
|
|
| |
| if args.prompt: |
| prompts = [args.prompt] |
| elif args.prompts: |
| with open(args.prompts) as f: |
| prompts = [l.strip() for l in f if l.strip()] |
| else: |
| with open(DEFAULT_PROMPTS_FILE) as f: |
| prompts = [l.strip() for l in f if l.strip()] |
|
|
| if not prompts: |
| print("No prompts found") |
| return |
|
|
| |
| pipe = SdxsPipeline.from_pretrained( |
| os.path.dirname(os.path.abspath(__file__)), |
| torch_dtype=torch.bfloat16, |
| trust_remote_code=True, |
| ).to("cuda:0") |
|
|
| 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, |
| show_progress_bar=False, |
| )[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}.jpg") |
| image.save(path, quality=95) |
| print(f" saved {path}") |
|
|
| |
| os.makedirs(MEDIA_DIR, exist_ok=True) |
| grid_path = os.path.join(MEDIA_DIR, "result_grid.jpg") |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| cols = min(4, 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_path, bbox_inches="tight", dpi=150, format="jpeg") |
| print(f"grid -> {grid_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
