Update app.py

app.py

@@ -23,7 +23,6 @@ subprocess.run('pip install flash-attn --no-build-isolation',
 # Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
 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
@@ -31,10 +30,6 @@ os.environ["HF_HOME"] = cache_path
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
-# Create gallery directory
-if not path.exists(gallery_path):
-    os.makedirs(gallery_path, exist_ok=True)
-
 # Florence 모델 초기화
 florence_models = {
     'gokaygokay/Florence-2-Flux-Large': AutoModelForCausalLM.from_pretrained(
@@ -63,12 +58,6 @@ def filter_prompt(prompt):
     "sex"    
     ]
     
-#        "nude", "naked", "nsfw", "porn", "sex", "explicit", "adult", "xxx",
-#        "erotic", "sensual", "seductive", "provocative", "intimate",
-#        "violence", "gore", "blood", "death", "kill", "murder", "torture",
-#        "drug", "suicide", "abuse", "hate", "discrimination"
-#    ]
-    
     prompt_lower = prompt.lower()
     
     for keyword in inappropriate_keywords:
@@ -107,59 +96,6 @@ pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretraine
     "CompVis/stable-diffusion-safety-checker"
 )
 
-
-
-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():
-    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 []
-
 @spaces.GPU
 def generate_caption(image, model_name='gokaygokay/Florence-2-Flux-Large'):
     image = Image.fromarray(image)
@@ -189,7 +125,7 @@ def process_and_save_image(height, width, steps, scales, prompt, seed):
     is_safe, filtered_prompt = filter_prompt(prompt)
     if not is_safe:
         gr.Warning("The prompt contains inappropriate content.")
-        return None, load_gallery()
+        return None
             
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
         try:
@@ -203,14 +139,10 @@ def process_and_save_image(height, width, steps, scales, prompt, seed):
                 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 get_random_seed():
     return torch.randint(0, 1000000, (1,)).item()
@@ -342,7 +274,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 "3.jpg",                
                 "1.jpg",
                 "4.jpg",
-              
             ]
             gr.Examples(
                 examples=example_images,
@@ -429,18 +360,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 label="Generated Image",
                 elem_classes=["output-image"]
             )
-            
-            gallery = gr.Gallery(
-                label="Generated Images Gallery",
-                show_label=True,
-                columns=[4],
-                rows=[2],
-                height="auto",
-                object_fit="cover",
-                elem_classes=["gallery-container"]
-            )
-            
-            gallery.value = load_gallery()
 
     # Event handlers
     caption_button.click(
@@ -452,7 +371,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     generate_btn.click(
         process_and_save_image,
         inputs=[height, width, steps, scales, prompt, seed],
-        outputs=[output, gallery]
+        outputs=[output]
     )
     
     randomize_seed.click(