Update app.py

app.py

@@ -8,10 +8,33 @@ from diffusers import ControlNetModel, StableDiffusionXLPipeline, AutoencoderKL
 from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
 import torch
 import os
-import uuid
+import time
+import glob
+
+# 一時ファイルの保存ディレクトリ
+TEMP_DIR = "temp_images"
+# 一時ファイルの保持期間（秒）
+FILE_RETENTION_PERIOD = 3600  # 1時間
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# 一時ディレクトリの作成
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+def cleanup_old_files():
+    """古い一時ファイルを削除する"""
+    current_time = time.time()
+    pattern = os.path.join(TEMP_DIR, "output_*.png")
+    
+    for file_path in glob.glob(pattern):
+        try:
+            # ファイルの最終更新時刻を取得
+            file_modified_time = os.path.getmtime(file_path)
+            if current_time - file_modified_time > FILE_RETENTION_PERIOD:
+                os.remove(file_path)
+        except Exception as e:
+            print(f"Error while cleaning up file {file_path}: {e}")
+
 pipe = StableDiffusionXLPipeline.from_single_file(
     "https://huggingface.co/Laxhar/noob_sdxl_beta/noob_hercules3/checkpoint/checkpoint-e2_s109089.safetensors/checkpoint-e2_s109089.safetensors",
     use_safetensors=True,
@@ -25,12 +48,15 @@ MAX_IMAGE_SIZE = 1216
 
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    # 古い一時ファイルの削除
+    cleanup_old_files()
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
+    # 画像生成
     output_image = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -41,10 +67,16 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         generator=generator
     ).images[0]
 
-    # PNGに変換する
-    output_image = output_image.convert("RGBA")
-
-    return output_image  # 画像オブジェクトを直接返す
+    # RGBモードで保存
+    if output_image.mode != 'RGB':
+        output_image = output_image.convert('RGB')
+    
+    # 一時ファイルとして保存
+    timestamp = int(time.time())
+    temp_filename = os.path.join(TEMP_DIR, f"output_{timestamp}.png")
+    output_image.save(temp_filename)
+    
+    return temp_filename
 
 css = """
 #col-container {
@@ -54,12 +86,12 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""
          Text-to-Image Demo
         using [Noob SDXL beta model](https://huggingface.co/Laxhar)
         """)
+        
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -68,21 +100,16 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-
             run_button = gr.Button("Run", scale=0)
 
-        # PNG形式でダウンロードできるようにfile_formatを指定
         result = gr.Image(
             label="Result",
             show_label=False,
-            type="pil",
-            elem_id="output_image",
-            show_download_button=True,
-            download_filename="output.png"
+            type="filepath",
+            elem_id="output_image"
         )
         
         with gr.Accordion("Advanced Settings", open=False):
-
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -140,4 +167,7 @@ with gr.Blocks(css=css) as demo:
         outputs=[result]
     )
 
-demo.queue().launch()
+# 起動時に古いファイルを削除
+cleanup_old_files()
+
+demo.queue().launch()