Update app.py

app.py

@@ -7,6 +7,9 @@ import random
 from diffusers import ControlNetModel, StableDiffusionXLPipeline, AutoencoderKL
 import cv2
 import torch
+import os
+import time
+import glob
 
 from diffusers import (
     DDIMScheduler,
@@ -20,21 +23,27 @@ from diffusers import (
     UniPCMultistepScheduler,
 )
 
+# 一時ファイルの管理設定
+TEMP_DIR = "temp_images"
+FILE_RETENTION_PERIOD = 3600  # 1時間
+os.makedirs(TEMP_DIR, exist_ok=True)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+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}")
 
-#vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-#pipe = StableDiffusionXLPipeline.from_pretrained(
-#    #"yodayo-ai/clandestine-xl-1.0", 
-#    torch_dtype=torch.float16, 
-#    use_safetensors=True,
-#    custom_pipeline="lpw_stable_diffusion_xl",
-#    add_watermarker=False #,
-#    #variant="fp16"
-#)
 pipe = StableDiffusionXLPipeline.from_single_file(
-    #"https://huggingface.co/Laxhar/noob_sdxl_beta/noob_hercules4/fp16/checkpoint-e0_s10000.safetensors/checkpoint-e0_s10000.safetensors",
     "https://huggingface.co/bluepen5805/illustrious_pencil-XL/illustrious_pencil-XL-v1.2.1.safetensors",
     use_safetensors=True,
     torch_dtype=torch.float16,
@@ -45,9 +54,10 @@ pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1216
 
-    
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, sampler_name):
+    # 古い一時ファイルの削除
+    cleanup_old_files()
 
     # サンプラーの設定
     if sampler_name == "DDIM":
@@ -71,7 +81,6 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     else:
         pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -87,11 +96,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 {
@@ -101,14 +115,12 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""
          Text-to-Image Demo
         using [illustrious_pencil-XL](https://huggingface.co/bluepen5805/illustrious_pencil-XL)
         """)
-        #yodayo-ai/clandestine-xl-1.0　
-        #yodayo-ai/holodayo-xl-2.1
+
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -117,20 +129,16 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-
             run_button = gr.Button("Run", scale=0)
 
         result = gr.Image(
             label="Result",
             show_label=False,
-            type="pil",
-            elem_id="output_image",
-            show_download_button=True#,
-            #download_filename="output.png"
+            type="filepath",  # filepathに変更
+            elem_id="output_image"
         )
         
         with gr.Accordion("Advanced Settings", open=False):
-
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -160,7 +168,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,#832,
+                    value=1024,
                 )
 
                 height = gr.Slider(
@@ -168,7 +176,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,#1216,
+                    value=1024,
                 )
 
             with gr.Row():
@@ -188,12 +196,13 @@ with gr.Blocks(css=css) as demo:
                     value=28,
                 )
 
-                
-
-    run_button.click(#lambda x: None, inputs=None, outputs=result).then(
+    run_button.click(
         fn=infer,
         inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, sampler_name],
         outputs=[result]
     )
 
+# 起動時に古いファイルを削除
+cleanup_old_files()
+
 demo.queue().launch()