File size: 3,556 Bytes
89a0dfc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """
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()
# collect prompts
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
# load pipeline
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}")
# save grid
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()
|