Update app.py

app.py

@@ -2,52 +2,111 @@ import spaces
 import gradio as gr
 import re
 from PIL import Image
+import io
 import os
-import gc  # Add garbage collection
-import psutil  # Add for memory monitoring
-
-# Set memory optimization flags
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:64"  # Reduced value
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Ensure using single GPU
-
 import numpy as np
 import torch
 from diffusers import FluxImg2ImgPipeline
+import tempfile
+import secrets
+import uuid
+import shutil
+import ssl
+from cryptography.fernet import Fernet
+import base64
+import hashlib
+import time
+import threading
 
-# Use float16 instead of bfloat16 for T4 compatibility
-dtype = torch.float16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Global encryption key for this session
+ENCRYPTION_KEY = Fernet.generate_key()
+cipher_suite = Fernet(ENCRYPTION_KEY)
 
-# Initialize the pipe directly during startup
-print("Loading model during startup...")
-torch.cuda.empty_cache()
-gc.collect()  # Force garbage collection
+# Configure SSL context for secure connections
+ssl_context = ssl.create_default_context()
+ssl_context.options |= ssl.OP_NO_TLSv1 | ssl.OP_NO_TLSv1_1  # Disable older TLS protocols
+ssl_context.set_ciphers('ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM:DHE+CHACHA20')
+ssl_context.check_hostname = True
+ssl_context.verify_mode = ssl.CERT_REQUIRED
 
-pipe = FluxImg2ImgPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell", 
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-    use_safetensors=True
-)
+# Secure temporary directory manager
+class SecureTempManager:
+    def __init__(self):
+        self.temp_dir = tempfile.mkdtemp(prefix='flux_secure_')
+        self.cleanup_timeout = 30  # seconds
+        self.file_registry = {}
+        
+    def get_secure_path(self, prefix="img"):
+        """Generate a secure random filename in the temp directory"""
+        filename = f"{prefix}_{uuid.uuid4().hex}_{secrets.token_hex(8)}.png"
+        filepath = os.path.join(self.temp_dir, filename)
+        
+        # Register file for cleanup
+        self.file_registry[filepath] = time.time()
+        
+        return filepath
+    
+    def cleanup_old_files(self):
+        """Clean up files older than the timeout"""
+        current_time = time.time()
+        for filepath, created_time in list(self.file_registry.items()):
+            if current_time - created_time > self.cleanup_timeout:
+                try:
+                    if os.path.exists(filepath):
+                        # Securely delete by overwriting with random data
+                        file_size = os.path.getsize(filepath)
+                        with open(filepath, 'wb') as f:
+                            f.write(os.urandom(file_size))
+                        # Then delete
+                        os.remove(filepath)
+                    # Remove from registry
+                    del self.file_registry[filepath]
+                except Exception as e:
+                    print(f"Error cleaning up file {filepath}: {e}")
+    
+    def cleanup_all(self):
+        """Clean up all files and remove temp directory"""
+        # First clean up individual files
+        for filepath in list(self.file_registry.keys()):
+            try:
+                if os.path.exists(filepath):
+                    os.remove(filepath)
+                del self.file_registry[filepath]
+            except:
+                pass
+        
+        # Then remove the directory
+        try:
+            if os.path.exists(self.temp_dir):
+                shutil.rmtree(self.temp_dir)
+        except:
+            pass
 
-# Enable attention slicing to reduce memory footprint
-pipe.enable_attention_slicing(1)
+# Initialize secure temp manager
+secure_temp = SecureTempManager()
 
-# Move to device immediately
-if torch.cuda.is_available():
-    pipe = pipe.to("cuda:0")
-else:
-    pipe = pipe.to("cpu")
+# Start a thread to periodically clean up old files
+def cleanup_thread_function():
+    while True:
+        secure_temp.cleanup_old_files()
+        time.sleep(5)  # Check every 5 seconds
 
-print("Model loaded successfully")
+cleanup_thread = threading.Thread(target=cleanup_thread_function)
+cleanup_thread.daemon = True  # Thread will exit when main program exits
+cleanup_thread.start()
+
+# Initialize model with proper settings
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(device)
 
 def sanitize_prompt(prompt):
-  # Allow only alphanumeric characters, spaces, and basic punctuation
-  allowed_chars = re.compile(r"[^a-zA-Z0-9\s.,!?-]")
-  sanitized_prompt = allowed_chars.sub("", prompt)
-  return sanitized_prompt
+    # Allow only alphanumeric characters, spaces, and basic punctuation
+    allowed_chars = re.compile(r"[^a-zA-Z0-9\s.,!?-]")
+    sanitized_prompt = allowed_chars.sub("", prompt)
+    return sanitized_prompt
 
-def convert_to_fit_size(original_width_and_height, maximum_size = 1024):
+def convert_to_fit_size(original_width_and_height, maximum_size = 2048):
     width, height = original_width_and_height
     if width <= maximum_size and height <= maximum_size:
         return width, height
@@ -66,64 +125,72 @@ def adjust_to_multiple_of_32(width: int, height: int):
     height = height - (height % 32)
     return width, height
 
-def resize_image(image: Image.Image, max_dim: int = 384) -> Image.Image:
-    """Resizes image to fit within max_dim while preserving aspect ratio"""
-    w, h = image.size
-    ratio = min(max_dim / w, max_dim / h)
-    if ratio < 1.0:
-        new_w = int(w * ratio)
-        new_h = int(h * ratio)
-        image = image.resize((new_w, new_h), Image.LANCZOS)
-    return image
-
-# Increase the timeout to 4 minutes
-@spaces.GPU(duration=740)
-def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step=2, progress=gr.Progress(track_tqdm=True)):
-    progress(0, desc="Starting and freeing memory")
-    # Free memory before processing
-    torch.cuda.empty_cache()
-    gc.collect()
+# Function to securely handle image data
+def secure_image_handler(image):
+    """Process image securely without exposing it to the file system"""
+    if image is None:
+        return None
+        
+    # If the image is already a PIL Image, use it directly
+    if isinstance(image, Image.Image):
+        return image
+        
+    # Otherwise, assume it's a file path or binary data
+    try:
+        if isinstance(image, str) and os.path.exists(image):
+            # It's a file path, load it securely
+            with open(image, 'rb') as f:
+                img_data = f.read()
+            
+            # Immediately delete the original file if it's in our temp directory
+            if image.startswith(secure_temp.temp_dir):
+                try:
+                    os.remove(image)
+                except:
+                    pass
+                
+            # Create image from binary data
+            return Image.open(io.BytesIO(img_data))
+        elif isinstance(image, bytes):
+            # It's binary data
+            return Image.open(io.BytesIO(image))
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return None
+
+@spaces.GPU(duration=120)
+def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step=4, progress=gr.Progress(track_tqdm=True)):
+    progress(0, desc="Starting")
     
-    progress(15, desc="Processing")
+    # Sanitize input
+    prompt = sanitize_prompt(prompt)
     
     def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=4):
+        # Secure image handling
+        image = secure_image_handler(image)
+        
         if image is None:
-            print("empty input image returned")
+            print("Empty input image returned")
             return None
             
-        # Convert webp to PNG if needed
-        if hasattr(image, 'format') and image.format == 'WEBP':
-            print("Converting WEBP to PNG")
-            rgb_im = image.convert('RGB')
-            image = rgb_im
-            
-        # Resize image to reduce memory usage
-        image = resize_image(image, max_dim=512)
-        
         generator = torch.Generator(device).manual_seed(seed)
-        fit_width, fit_height = convert_to_fit_size(image.size, maximum_size=512)
+        fit_width, fit_height = convert_to_fit_size(image.size)
         width, height = adjust_to_multiple_of_32(fit_width, fit_height)
         image = image.resize((width, height), Image.LANCZOS)
-        
-        progress(30, desc="Processing image")
-        
-        # Use autocast for better memory efficiency
-        with torch.cuda.amp.autocast(dtype=torch.float16):
-            with torch.no_grad():
-                output = pipe(
-                    prompt=prompt, 
-                    image=image, 
-                    generator=generator, 
-                    strength=strength, 
-                    width=width, 
-                    height=height,
-                    guidance_scale=0, 
-                    num_inference_steps=num_inference_steps, 
-                    max_sequence_length=256
-                )
-        
-        progress(90, desc="Finalizing")
-        
+    
+        # Process the image
+        output = pipe(
+            prompt=prompt, 
+            image=image, 
+            generator=generator, 
+            strength=strength, 
+            width=width, 
+            height=height,
+            guidance_scale=0, 
+            num_inference_steps=num_inference_steps, 
+            max_sequence_length=256
+        )
+    
         pil_image = output.images[0]
         new_width, new_height = pil_image.size
     
@@ -132,9 +199,20 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
             return resized_image
         return pil_image
     
+    # Process the image
     output = process_img2img(image, prompt, strength, seed, inference_step)
-    progress(100, desc="Done")
-    return output
+    
+    # Instead of returning image directly, save to secure temp location and return path
+    if output is not None:
+        # Convert output to secure in-memory format
+        output_buffer = io.BytesIO()
+        output.save(output_buffer, format="PNG")
+        output_buffer.seek(0)
+        
+        # Return the image data directly (Gradio will handle it)
+        return output
+    
+    return None
 
 def read_file(path: str) -> str:
     with open(path, 'r', encoding='utf-8') as f:
@@ -168,48 +246,102 @@ css="""
 }
 """
 
+# Custom HTTP headers for security
+custom_headers = {
+    "Strict-Transport-Security": "max-age=63072000; includeSubDomains; preload",
+    "X-Content-Type-Options": "nosniff",
+    "X-Frame-Options": "SAMEORIGIN",
+    "Content-Security-Policy": "default-src 'self'; img-src 'self' data:; style-src 'self' 'unsafe-inline';",
+    "Referrer-Policy": "strict-origin-when-cross-origin",
+    "Permissions-Policy": "camera=(), microphone=(), geolocation=()",
+    "Cache-Control": "no-store, max-age=0"
+}
+
+# Create Gradio app with enhanced security
 with gr.Blocks(css=css, elem_id="demo-container") as demo:
     with gr.Column():
         gr.HTML(read_file("demo_header.html"))
         gr.HTML(read_file("demo_tools.html"))
     with gr.Row():
-                with gr.Column():
-                    image = gr.Image(height=800,sources=['upload','clipboard'],image_mode='RGB', elem_id="image_upload", type="pil", label="Upload")
-                    with gr.Row(elem_id="prompt-container",  equal_height=False):
-                        with gr.Row():
-                            prompt = gr.Textbox(label="Prompt",value="a women",placeholder="Your prompt (what you want in place of what is erased)", elem_id="prompt")
-                            
-                    btn = gr.Button("Img2Img", elem_id="run_button",variant="primary")
+        with gr.Column():
+            image = gr.Image(
+                height=800,
+                sources=['upload','clipboard'],
+                image_mode='RGB', 
+                elem_id="image_upload", 
+                type="pil", 
+                label="Upload"
+            )
+            with gr.Row(elem_id="prompt-container", equal_height=False):
+                with gr.Row():
+                    prompt = gr.Textbox(
+                        label="Prompt",
+                        value="a women",
+                        placeholder="Your prompt (what you want in place of what is erased)", 
+                        elem_id="prompt"
+                    )
                     
-                with gr.Accordion(label="Advanced Settings", open=False):
-                        with gr.Row( equal_height=True):
-                            strength = gr.Number(value=0.65, minimum=0, maximum=0.75, step=0.01, label="strength")
-                            seed = gr.Number(value=100, minimum=0, step=1, label="seed")
-                            inference_step = gr.Number(value=2, minimum=1, maximum=4, step=1, label="inference_step")
-                        id_input=gr.Text(label="Name", visible=False)
-                            
-                with gr.Column():
-                    image_out = gr.Image(height=800,sources=[],label="Output", elem_id="output-img",format="jpg")
-                   
+            btn = gr.Button("Img2Img", elem_id="run_button", variant="primary")
+            
+            with gr.Accordion(label="Advanced Settings", open=False):
+                with gr.Row(equal_height=True):
+                    strength = gr.Number(value=0.75, minimum=0, maximum=0.75, step=0.01, label="strength")
+                    seed = gr.Number(value=100, minimum=0, step=1, label="seed")
+                    inference_step = gr.Number(value=4, minimum=1, step=4, label="inference_step")
+                id_input=gr.Text(label="Name", visible=False)
+                    
+        with gr.Column():
+            image_out = gr.Image(
+                height=800,
+                sources=[],
+                label="Output", 
+                elem_id="output-img",
+                format="jpg"
+            )
+
     gr.Examples(
-               examples=[
-                    ["examples/draw_input.jpg", "examples/draw_output.jpg","a women ,eyes closed,mouth opened"],
-                    ["examples/draw-gimp_input.jpg", "examples/draw-gimp_output.jpg","a women ,eyes closed,mouth opened"],
-                    ["examples/gimp_input.jpg", "examples/gimp_output.jpg","a women ,hand on neck"],
-                    ["examples/inpaint_input.jpg", "examples/inpaint_output.jpg","a women ,hand on neck"]
-                         ]
-,
-                inputs=[image,image_out,prompt],
+        examples=[
+            ["examples/draw_input.jpg", "examples/draw_output.jpg", "a women ,eyes closed,mouth opened"],
+            ["examples/draw-gimp_input.jpg", "examples/draw-gimp_output.jpg", "a women ,eyes closed,mouth opened"],
+            ["examples/gimp_input.jpg", "examples/gimp_output.jpg", "a women ,hand on neck"],
+            ["examples/inpaint_input.jpg", "examples/inpaint_output.jpg", "a women ,hand on neck"]
+        ],
+        inputs=[image, image_out, prompt],
     )
     gr.HTML(
-       gr.HTML(read_file("demo_footer.html"))
+        gr.HTML(read_file("demo_footer.html"))
     )
     gr.on(
         triggers=[btn.click, prompt.submit],
-        fn = process_images,
-        inputs = [image,prompt,strength,seed,inference_step],
-        outputs = [image_out]
+        fn=process_images,
+        inputs=[image, prompt, strength, seed, inference_step],
+        outputs=[image_out]
     )
 
+# Register shutdown handler to clean up
+import atexit
+atexit.register(secure_temp.cleanup_all)
+
 if __name__ == "__main__":
-    demo.launch(share=True, show_error=True)
+    # Launch with security settings
+    demo.launch(
+        share=True, 
+        show_error=True,
+        ssl_verify=True,
+        ssl_certfile="server.crt",  # You'll need to generate these
+        ssl_keyfile="server.key",   # for production use
+        ssl_keyfile_password=None,
+        ssl_client_cert_chain="chain.pem",  # Optional for client cert verification
+        favicon_path=None, 
+        server_name="0.0.0.0",  # Listen on all interfaces
+        server_port=7860,       # Default Gradio port
+        inbrowser=False,       
+        debug=False,           # Disable in production
+        quiet=True,             # Less logging for security
+        height=900,
+        width=1600,
+        enable_queue=True,
+        max_threads=20,
+        auth=None,              # Enable if you need authentication
+        root_path=""
+    )