Browse Source

change run_cli.py to support save prompts.

pull/100/head
shiyiming 2 years ago
parent
commit
8940096837
  1. 23
      run_cli.py

23
run_cli.py

@ -1,3 +1,5 @@
# python run_cli.py -f /data/hdd1/shiym/competition/GitHubClone/diffusers/examples/custom_diffusion/real_reg/samples_girl/images -c ViT-L-14/openai -d cuda -m fast
#!/usr/bin/env python3 #!/usr/bin/env python3
import argparse import argparse
import csv import csv
@ -81,15 +83,20 @@ def main():
prompts.append(prompt) prompts.append(prompt)
print(prompt) print(prompt)
if len(prompts): # 写入到本地文件 caption.txt 中
csv_path = os.path.join(args.folder, 'desc.csv') with open("/data/hdd1/shiym/competition/GitHubClone/diffusers/examples/custom_diffusion/real_reg/samples_girl/caption.txt", "w", encoding="utf-8") as f:
with open(csv_path, 'w', encoding='utf-8', newline='') as f: for prompt in prompts:
w = csv.writer(f, quoting=csv.QUOTE_MINIMAL) f.write(prompt + "\n")
w.writerow(['image', 'prompt'])
for file, prompt in zip(files, prompts): # if len(prompts):
w.writerow([file, prompt]) # csv_path = os.path.join(args.folder, 'desc.csv')
# with open(csv_path, 'w', encoding='utf-8', newline='') as f:
# w = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
# w.writerow(['image', 'prompt'])
# for file, prompt in zip(files, prompts):
# w.writerow([file, prompt])
print(f"\n\n\n\nGenerated {len(prompts)} and saved to {csv_path}, enjoy!") # print(f"\n\n\n\nGenerated {len(prompts)} and saved to {csv_path}, enjoy!")
if __name__ == "__main__": if __name__ == "__main__":
main() main()

Loading…
Cancel
Save