Update image saving logic to correctly handle prompts and save the final image

This commit is contained in:
2026-05-07 11:05:05 +02:00
parent e32f6526f6
commit ad68d6bb72

View File

@@ -12,7 +12,6 @@ import json
import math
import os
import re
import shutil
import tempfile
import torch
from PIL import Image, ImageDraw
@@ -555,7 +554,9 @@ def generate(
except Exception as e: # noqa: BLE001
print(f"Upscaling failed ({upscale}); keeping seamless image: {e}")
shutil.move(final_path, output_path)
with Image.open(final_path) as final_img:
final_img.load()
save_png_with_prompt(final_img, output_path, prompt)
try:
with Image.open(output_path) as _im:
print(f"→ Final image written to {output_path} [{_im.size[0]}x{_im.size[1]}]")