Update app.py

app.py

@@ -15,9 +15,6 @@ from PIL import Image
 
 # Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-# Use PERSISTENT_DIR environment variable for Spaces
-PERSISTENT_DIR = os.environ.get("PERSISTENT_DIR", ".")
-gallery_path = path.join(PERSISTENT_DIR, "gallery")
 
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
@@ -25,10 +22,6 @@ os.environ["HF_HOME"] = cache_path
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
-# Create gallery directory if it doesn't exist
-if not path.exists(gallery_path):
-    os.makedirs(gallery_path, exist_ok=True)
-
 def filter_prompt(prompt):
     # 부적절한 키워드 목록
     inappropriate_keywords = [
@@ -101,45 +94,6 @@ footer {display: none !important}
     -webkit-background-clip: text;
     -webkit-text-fill-color: transparent;
 }
-#gallery {
-    width: 100% !important;
-    max-width: 100% !important;
-    overflow: visible !important;
-}
-#gallery > div {
-    width: 100% !important;
-    max-width: none !important;
-}
-#gallery > div > div {
-    width: 100% !important;
-    display: grid !important;
-    grid-template-columns: repeat(5, 1fr) !important;
-    gap: 16px !important;
-    padding: 16px !important;
-}
-.gallery-container {
-    background: rgba(255, 255, 255, 0.05);
-    border-radius: 8px;
-    margin-top: 10px;
-    width: 100% !important;
-    box-sizing: border-box !important;
-}
-.gallery-item {
-    width: 100% !important;
-    aspect-ratio: 1 !important;
-    overflow: hidden !important;
-    border-radius: 4px !important;
-}
-.gallery-item img {
-    width: 100% !important;
-    height: 100% !important;
-    object-fit: cover !important;
-    border-radius: 4px !important;
-    transition: transform 0.2s;
-}
-.gallery-item img:hover {
-    transform: scale(1.05);
-}
 .output-image {
     width: 100% !important;
     max-width: 100% !important;
@@ -152,76 +106,8 @@ footer {display: none !important}
     width: 100% !important;
     max-width: 100% !important;
 }
-.gallery-container::-webkit-scrollbar {
-    display: none !important;
-}
-.gallery-container {
-    -ms-overflow-style: none !important;
-    scrollbar-width: none !important;
-}
-#gallery > div {
-    width: 100% !important;
-    max-width: 100% !important;
-}
-#gallery > div > div {
-    width: 100% !important;
-    max-width: 100% !important;
-}
 """
 
-def save_image(image):
-    """Save the generated image and return the path"""
-    try:
-        if not os.path.exists(gallery_path):
-            try:
-                os.makedirs(gallery_path, exist_ok=True)
-            except Exception as e:
-                print(f"Failed to create gallery directory: {str(e)}")
-                return None
-        
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        random_suffix = os.urandom(4).hex()
-        filename = f"generated_{timestamp}_{random_suffix}.png"
-        filepath = os.path.join(gallery_path, filename)
-        
-        try:
-            if isinstance(image, Image.Image):
-                image.save(filepath, "PNG", quality=100)
-            else:
-                image = Image.fromarray(image)
-                image.save(filepath, "PNG", quality=100)
-            
-            if not os.path.exists(filepath):
-                print(f"Warning: Failed to verify saved image at {filepath}")
-                return None
-                
-            return filepath
-        except Exception as e:
-            print(f"Failed to save image: {str(e)}")
-            return None
-            
-    except Exception as e:
-        print(f"Error in save_image: {str(e)}")
-        return None
-
-def load_gallery():
-    """Load all images from the gallery directory"""
-    try:
-        os.makedirs(gallery_path, exist_ok=True)
-        
-        image_files = []
-        for f in os.listdir(gallery_path):
-            if f.lower().endswith(('.png', '.jpg', '.jpeg')):
-                full_path = os.path.join(gallery_path, f)
-                image_files.append((full_path, os.path.getmtime(full_path)))
-        
-        image_files.sort(key=lambda x: x[1], reverse=True)
-        
-        return [f[0] for f in image_files]
-    except Exception as e:
-        print(f"Error loading gallery: {str(e)}")
-        return []
-
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.HTML('<div class="title">AI Image Generator</div>')
@@ -322,27 +208,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 elem_id="output-image",
                 elem_classes=["output-image", "fixed-width"]
             )
-            
-            gallery = gr.Gallery(
-                label="Generated Images Gallery",
-                show_label=True,
-                elem_id="gallery",
-                columns=[4],
-                rows=[2],
-                height="auto",
-                object_fit="cover",
-                elem_classes=["gallery-container", "fixed-width"]
-            )
-            
-            gallery.value = load_gallery()
     
     @spaces.GPU
-    def process_and_save_image(height, width, steps, scales, prompt, seed):
+    def process_image(height, width, steps, scales, prompt, seed):
         # 프롬프트 필터링
         is_safe, filtered_prompt = filter_prompt(prompt)
         if not is_safe:
             gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
-            return None, load_gallery()
+            return None
             
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
             try:
@@ -356,22 +229,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     max_sequence_length=256
                 ).images[0]
                 
-                saved_path = save_image(generated_image)
-                if saved_path is None:
-                    print("Warning: Failed to save generated image")
-                
-                return generated_image, load_gallery()
+                return generated_image
             except Exception as e:
                 print(f"Error in image generation: {str(e)}")
-                return None, load_gallery()
+                return None
     
     def update_seed():
         return get_random_seed()
 
     generate_btn.click(
-        process_and_save_image,
+        process_image,
         inputs=[height, width, steps, scales, prompt, seed],
-        outputs=[output, gallery]
+        outputs=[output]
     )
     
     randomize_seed.click(
@@ -385,4 +254,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(allowed_paths=[PERSISTENT_DIR])
+    demo.launch()