From 8f7ef6dfce96979dc4c9a11350ef7703b843d9df Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 28 Jan 2024 20:50:22 -0600 Subject: [PATCH] remove references to /tmp --- api/onnx_web/diffusers/pipelines/highres.py | 7 ++++--- api/onnx_web/output.py | 1 + 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/api/onnx_web/diffusers/pipelines/highres.py b/api/onnx_web/diffusers/pipelines/highres.py index 39112b4c..96979291 100644 --- a/api/onnx_web/diffusers/pipelines/highres.py +++ b/api/onnx_web/diffusers/pipelines/highres.py @@ -362,7 +362,7 @@ def export_unet(pipeline, output_path, unet_sample_size=1024): ) -def load_and_export(source="stabilityai/sd-x2-latent-upscaler"): +def load_and_export(output, source="stabilityai/sd-x2-latent-upscaler"): from pathlib import Path from diffusers import StableDiffusionLatentUpscalePipeline @@ -370,11 +370,12 @@ def load_and_export(source="stabilityai/sd-x2-latent-upscaler"): upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained( source, torch_dtype=torch.float32 ) - export_unet(upscaler, Path("/tmp/latent-upscaler")) + export_unet(upscaler, Path(output)) def load_and_run( prompt, + output, source="stabilityai/sd-x2-latent-upscaler", checkpoint="../models/stable-diffusion-onnx-v1-5", ): @@ -396,7 +397,7 @@ def load_and_run( # run result = highres.text2img(prompt, num_inference_steps=25, num_upscale_steps=25) image = result.images[0] - image.save("/tmp/highres.png") + image.save(output) class RetorchModel: diff --git a/api/onnx_web/output.py b/api/onnx_web/output.py index 2aa221a9..57049132 100644 --- a/api/onnx_web/output.py +++ b/api/onnx_web/output.py @@ -95,6 +95,7 @@ def save_result( thumbnails = [] for image, filename in zip(images, result.thumbnails): + # TODO: only make a thumbnail if the image is larger than the thumbnail size thumbnail = image.copy() thumbnail.thumbnail((server.thumbnail_size, server.thumbnail_size))