diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 3b0a0d6a..af0e9b33 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -1227,6 +1227,58 @@ class GenerateImageConfig: # join with folder return os.path.join(self.output_folder, filename) + def save_image_atomic(self, image, count: int = 0, max_count=0): + # write into a hidden tmp subfolder, then atomically move into place so + # watchers (UI/CDN) never see and cache a partially written file. Wraps + # self.save_image so it also covers models that replace that function. + real_folder = self.output_folder + tmp_folder = os.path.join(real_folder, '.tmp') + os.makedirs(tmp_folder, exist_ok=True) + self.output_folder = tmp_folder + try: + self.save_image(image, count, max_count) + finally: + self.output_folder = real_folder + files = os.listdir(tmp_folder) + # thumbs move into place first so they already exist when the media + # file appears in the samples folder + thumbs_folder = os.path.join(real_folder, '.thumbs') + for file in files: + tmp_thumb = os.path.join(tmp_folder, file + '.thumb') + try: + if self._generate_thumbnail(os.path.join(tmp_folder, file), tmp_thumb): + os.makedirs(thumbs_folder, exist_ok=True) + os.replace(tmp_thumb, os.path.join(thumbs_folder, file + '.jpg')) + except Exception as e: + print(f"Failed to generate thumbnail for {file}: {e}") + for file in files: + os.replace(os.path.join(tmp_folder, file), os.path.join(real_folder, file)) + + def _generate_thumbnail(self, media_path, thumb_path): + # 300x300 center-cropped 90% jpg. Returns True if one was written. + from PIL import Image as PILImage + ext = os.path.splitext(media_path)[1].lower() + img = None + if ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp']: + img = PILImage.open(media_path) # animated formats open on the first frame + elif ext == '.mp4': + import cv2 + cap = cv2.VideoCapture(media_path) + ok, frame = cap.read() + cap.release() + if ok: + img = PILImage.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + if img is None: + return False + img = img.convert('RGB') + w, h = img.size + side = min(w, h) + left = (w - side) // 2 + top = (h - side) // 2 + img = img.crop((left, top, left + side, top + side)).resize((300, 300), PILImage.LANCZOS) + img.save(thumb_path, format='JPEG', quality=90) + return True + def save_image(self, image, count: int = 0, max_count=0): # make parent dirs os.makedirs(self.output_folder, exist_ok=True) diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 795e6ae3..394adcfc 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -666,7 +666,7 @@ class BaseModel: extra, ) - gen_config.save_image(img, i) + gen_config.save_image_atomic(img, i) gen_config.log_image(img, i) self._after_sample_image(i, len(image_configs)) flush() diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index a51e9b55..7a6c7855 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -1715,7 +1715,7 @@ class StableDiffusion: generator=generator, ).images[0] - gen_config.save_image(img, i) + gen_config.save_image_atomic(img, i) gen_config.log_image(img, i) self._after_sample_image(i, len(image_configs)) flush() diff --git a/ui/cron/fileServer.ts b/ui/cron/fileServer.ts index 52b0f50a..02515470 100644 --- a/ui/cron/fileServer.ts +++ b/ui/cron/fileServer.ts @@ -118,7 +118,7 @@ async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, pr .map(decodeURIComponent) .join('/'); - const resolvedFilePath = path.resolve(decodedFilePath); + let resolvedFilePath = path.resolve(decodedFilePath); const roots = await getRoots(); const allowedDirs = isImg ? [roots.datasets, roots.training, roots.data] : [roots.datasets, roots.training]; const isAllowed = allowedDirs.some( @@ -131,6 +131,17 @@ async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, pr return; } + // ?thumb=1 serves the pre-generated 300x300 jpg from the sibling .thumbs + // folder (..jpg) when it exists; otherwise falls through to + // the full file exactly as before. Mirrors the Next.js /api/img route. + if (isImg && new URL(req.url || '', 'http://localhost').searchParams.has('thumb')) { + const thumbPath = path.join(path.dirname(resolvedFilePath), '.thumbs', path.basename(resolvedFilePath) + '.jpg'); + const thumbStat = await fs.promises.stat(thumbPath).catch(() => null); + if (thumbStat && thumbStat.isFile()) { + resolvedFilePath = thumbPath; + } + } + let stat: fs.Stats; try { stat = await fs.promises.stat(resolvedFilePath); diff --git a/ui/src/app/api/img/[...imagePath]/route.ts b/ui/src/app/api/img/[...imagePath]/route.ts index e41362cf..70c25ab3 100644 --- a/ui/src/app/api/img/[...imagePath]/route.ts +++ b/ui/src/app/api/img/[...imagePath]/route.ts @@ -45,7 +45,7 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat // Security check: resolve the path so any `..` segments are collapsed, // then ensure it's still under an allowed root. (Plain `.includes('..')` // false-positives on filenames that contain `..` as text, e.g. an ellipsis.) - const resolved = path.resolve(filepath); + let resolved = path.resolve(filepath); const isAllowed = allowedDirs.some( allowedDir => resolved === allowedDir || resolved.startsWith(allowedDir + path.sep), ); @@ -55,6 +55,17 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat return new NextResponse('Access denied', { status: 403 }); } + // ?thumb=1 serves the pre-generated 300x300 jpg from the sibling .thumbs + // folder (..jpg) when it exists; otherwise falls through to + // the full file exactly as before. + if (request.nextUrl.searchParams.has('thumb')) { + const thumbPath = path.join(path.dirname(resolved), '.thumbs', path.basename(resolved) + '.jpg'); + const thumbStat = await fs.promises.stat(thumbPath).catch(() => null); + if (thumbStat && thumbStat.isFile()) { + resolved = thumbPath; + } + } + // Bail out early if the client already gave up if (request.signal.aborted) { return new NextResponse(null, { status: 499 }); diff --git a/ui/src/app/api/jobs/[jobID]/samples/route.ts b/ui/src/app/api/jobs/[jobID]/samples/route.ts index 6a2aafe2..04c461ba 100644 --- a/ui/src/app/api/jobs/[jobID]/samples/route.ts +++ b/ui/src/app/api/jobs/[jobID]/samples/route.ts @@ -25,8 +25,12 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s return NextResponse.json({ samples: [] }); } - // find all img (png, jpg, jpeg) files in the samples folder - const samples = (await fs.promises.readdir(samplesFolder)) + // find all img (png, jpg, jpeg) files in the samples folder. Thumbnails + // live in the hidden .thumbs subfolder (and partial writes in .tmp) — the + // isFile check keeps those directories out even if their names ever match. + const samples = (await fs.promises.readdir(samplesFolder, { withFileTypes: true })) + .filter(entry => entry.isFile()) + .map(entry => entry.name) .filter(file => { return file.endsWith('.png') || file.endsWith('.jpg') || file.endsWith('.jpeg') || file.endsWith('.webp') || file.endsWith('.mp4') || file.endsWith('mp3') || file.endsWith('wav') || file.endsWith('flac') || file.endsWith('ogg'); }) diff --git a/ui/src/components/SampleImageCard.tsx b/ui/src/components/SampleImageCard.tsx index 0973e382..f27222b3 100644 --- a/ui/src/components/SampleImageCard.tsx +++ b/ui/src/components/SampleImageCard.tsx @@ -32,10 +32,11 @@ const SampleImageCard: React.FC = ({ const [isVisible, setIsVisible] = useState(false); const [loaded, setLoaded] = useState(false); const [blobUrl, setBlobUrl] = useState(null); + // videos with no pre-generated thumb (older samples) fall back to the