把一串提示词跑完,循环本身不是难点;撑爆显存的常见原因是循环里没断掉的引用、每轮重复加载或重建模型、以及为了省事把每轮生成的张量或图像对象留在变量或列表里。可行的顺序是:先把单张出图跑通并保存成文件,再用循环包住它,每轮结束做一次释放,同时把每轮的进度写进日志,中断后按日志跳过已完成的条目,最后再把输出目录和文件命名参数化,避免重跑互相覆盖。
批量出图的显存问题,通常不在 for 循环而在循环里留下的对象。先确认单张出图能保存到文件并在返回后显存回落到接近加载模型后的基线,再套循环;每轮在 finally 里解除引用、收集垃圾、调用显存缓存清理,并把状态逐行写入日志。如果显存随轮次单调上涨,说明还有引用或状态在累积,应先定位再谈批量大小和并发;释放调用只是收回未被引用的缓存,不能替代修复泄漏。
把提示词整理成可逐行读取的列表文件
用逐行 JSON 作为输入,比 CSV 或直接写死在脚本里更好维护:一行一条、可加注释、出错时能定位到具体行。空行和以 # 开头的行在读入时跳过,方便临时停用某几条提示词。
# tasks.jsonl
{"prompt": "a cat on a windowsill, soft light", "seed": 1234, "width": 1024, "height": 1024}
{"prompt": "城市夜景,雨后街道", "seed": 42, "width": 768, "height": 1024, "name": "city_rain"}
- prompt:必填,单行字符串,不要在里面混入换行,避免日志和文件名不好处理。
- seed:建议显式写;缺省时由脚本生成,但必须写进日志,否则续跑时无法复现同一张图。
- width、height:建议先用模型支持范围内的保守尺寸跑通,再按需放大。
- name:可选,用于文件名里的可读标识;不填时从 prompt 摘要生成。
读取位置建议单独写一个生成器,返回行号和解析后的对象。行号是比 prompt 字符串更稳定的任务标识,长提示词做去重键容易出错。
import json
def read_tasks(path):
with open(path, encoding='utf-8') as f:
for lineno, line in enumerate(f, 1):
line = line.strip()
if not line or line.startswith('#'):
continue
try:
yield lineno, json.loads(line)
except json.JSONDecodeError as e:
print(f'第 {lineno} 行解析失败: {e}') # 记日志后跳过,不要中断整批
先写单张出图函数并跑通一张
把单次推理封装成一个函数,入参只保留提示词、种子、尺寸和输出路径;模型或管道对象在函数外加载一次,通过参数传进来。这样函数的显存占用是“一份模型加一轮中间量”,不会随轮次叠加,也不会每轮重新加载权重。
import os, torch
def generate_one(pipe, prompt, seed, width, height, out_path):
gen = torch.Generator(device='cuda').manual_seed(seed)
# ---- 替换点:这里放 SeFi-Image 的实际调用方式 ----
image = pipe(prompt=prompt, width=width, height=height, generator=gen).images[0]
# ------------------------------------------------
os.makedirs(os.path.dirname(out_path), exist_ok=True)
image.save(out_path)
return out_path
验证方式很具体:先只调用一次,确认返回路径存在、图片能打开、尺寸与入参一致;再看调用前后 nvidia-smi 的显存占用。返回后占用不一定与调用前完全相等,出现一定缓存属于正常,但应当回落到接近“刚加载完模型”的基线,而不是停在推理峰值。如果单张跑完就明显不下来,先处理这一次的问题,不要继续加循环。
用循环包住它并在每轮后释放显存
循环骨架里要守住两条:循环内不做模型加载,循环内不把图像或张量对象累积到列表里,日志只保存路径字符串。释放动作放在每轮的 finally 中,异常和正常结束都执行。
pipe = load_pipeline() # 放在循环外,只加载一次
for lineno, task in read_tasks('tasks.jsonl'):
out_path = build_path(run_dir, lineno, task)
image = None
try:
image = generate_one(pipe, task['prompt'], task.get('seed', 0),
task.get('width', 1024), task.get('height', 1024),
out_path)
log(lineno, 'done', out_path)
except Exception as e:
log(lineno, 'failed', repr(e))
finally:
del image
gc.collect()
torch.cuda.empty_cache()
释放的位置值得说明:del 加 gc.collect() 是真正解除对上一轮中间结果的引用,torch.cuda.empty_cache() 只是把已经没人引用的缓存块交还给驱动,它不解决泄漏,也不应被当作提速手段。观察方法可以另开一个终端执行 nvidia-smi -l 1,重点看每轮结束后的占用是否回到同一条水平线;如果逐轮缓慢抬升且不回落,多半是还有引用或某种状态在累积,先定位再考虑并发或批量尺寸。如果单张出图本身就接近显存上限,能做的通常是降低尺寸、减少并发或分批跑,而不是靠释放调用硬撑。
记录每轮的输出路径与失败原因
日志按行追加写,一行一条 JSON,写完立即 flush。只在整批结束时汇总写一次文件,进程被中断时最后一段记录很容易丢。
{"run_id": "run-0821-a", "index": 7, "prompt_hash": "3f9a1c2e", "seed": 42,
"width": 768, "height": 1024, "status": "done",
"out_path": "out/run-0821-a/0007-city-rain-s42-768x1024.png", "elapsed_s": 12.4}
- 字段建议:run_id、index(列表文件行号)、prompt_hash(取摘要前 8 位即可,不必存全文)、seed、尺寸、status、out_path、error_type、error_msg(截断到一两百字符)、elapsed_s。
- 标记方式:开跑前先写一条 running,成功改为 done,异常写 failed。中断后能从日志里看出卡在哪一条,而不是只剩一个空目录。
- 续跑逻辑:启动时扫描日志,收集 status 为 done 的 index 集合,循环里遇到就跳过;failed 默认也跳过,用一个 `--retry-failed` 参数控制是否重试,避免一有坏条目就每轮重复报错。
- 如果换过参数或模型版本,建议换新的 run_id 重跑,而不是在旧目录上追加日志,否则“跳过”会基于过期的记录判断。
给脚本加参数化输出目录与文件命名
命名规则把 run_id、序号、可读标识、种子和尺寸拼在一起,信息足够从文件名反查任务,也不易撞名。序号取列表文件行号并补零,保证目录内按生成顺序排列。
out/
run-0821-a/
tasks.jsonl
run.log.jsonl
0001-a-cat-on-a-windowsill-s1234-1024x1024.png
0002-city-rain-s42-768x1024.png
可读标识从 name 字段取,没有就用 prompt 的 hash;slug 化时去掉空格和标点,限制长度。覆盖策略默认保守:目标文件已存在就写一条 skipped 日志并跳过,需要重画时显式加 `--overwrite`,或者直接换新的 run_id 生成新目录,把新旧结果分开保存。这样一批跑到一半中断,重跑只会补齐缺失的条目,不会把已经确认过的那部分重画一遍。