Update app.py

app.py

@@ -3,103 +3,69 @@ import gradio as gr
 import re
 from PIL import Image
 import io
+import base64
 import os
 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
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
-# Global encryption key for this session
-ENCRYPTION_KEY = Fernet.generate_key()
-cipher_suite = Fernet(ENCRYPTION_KEY)
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Configure SSL context for secure connections - FIXED TO AVOID DEPRECATION WARNING
-ssl_context = ssl.create_default_context()
-# Use the recommended modern approach instead of deprecated options
-ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2
-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.bfloat16).to(device)
 
-# 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
+# Encryption setup
+def generate_key(password, salt=None):
+    if salt is None:
+        salt = os.urandom(16)
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=100000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
+    return key, salt
+
+def encrypt_image(image, password="default_password"):
+    # Convert PIL Image to bytes
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='PNG')
+    img_byte_arr = img_byte_arr.getvalue()
     
-    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}")
+    # Generate key for encryption
+    key, salt = generate_key(password)
+    cipher = Fernet(key)
     
-    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
-
-# Initialize secure temp manager
-secure_temp = SecureTempManager()
-
-# 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
-
-cleanup_thread = threading.Thread(target=cleanup_thread_function)
-cleanup_thread.daemon = True  # Thread will exit when main program exits
-cleanup_thread.start()
+    # Encrypt the image bytes
+    encrypted_data = cipher.encrypt(img_byte_arr)
+    
+    # Return the encrypted data and salt (needed for decryption)
+    return {
+        'encrypted_data': base64.b64encode(encrypted_data).decode('utf-8'),
+        'salt': base64.b64encode(salt).decode('utf-8'),
+        'original_width': image.width,
+        'original_height': image.height
+    }
 
-# 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 decrypt_image(encrypted_data_dict, password="default_password"):
+    # Extract the encrypted data and salt
+    encrypted_data = base64.b64decode(encrypted_data_dict['encrypted_data'])
+    salt = base64.b64decode(encrypted_data_dict['salt'])
+    
+    # Regenerate the key using the provided salt
+    key, _ = generate_key(password, salt)
+    cipher = Fernet(key)
+    
+    # Decrypt the data
+    decrypted_data = cipher.decrypt(encrypted_data)
+    
+    # Convert bytes back to PIL Image
+    image = Image.open(io.BytesIO(decrypted_data))
+    return image
 
 def sanitize_prompt(prompt):
     # Allow only alphanumeric characters, spaces, and basic punctuation
@@ -126,71 +92,21 @@ def adjust_to_multiple_of_32(width: int, height: int):
     height = height - (height % 32)
     return width, height
 
-# 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)):
+def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step=4, encrypt_password="default_password", progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Starting")
-    
-    # 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
-            
         generator = torch.Generator(device).manual_seed(seed)
         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)
     
-        # 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
-        )
+        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
@@ -200,11 +116,24 @@ 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)
     
-    # Return the image directly (Gradio will handle it)
-    return output
+    # Encrypt the output image
+    if output is not None:
+        encrypted_output = encrypt_image(output, encrypt_password)
+        
+        # For display purposes, we'll create a placeholder image with text indicating encryption
+        placeholder = Image.new('RGB', (output.width, output.height), color=(220, 220, 220))
+        return {
+            "display_image": placeholder,
+            "encrypted_data": encrypted_output
+        }
+    return None
+
+def save_encrypted_image(encrypted_data, filename="encrypted_image.enc"):
+    with open(filename, 'w') as f:
+        json.dump(encrypted_data, f)
+    return f"Encrypted image saved as {filename}"
 
 def read_file(path: str) -> str:
     with open(path, 'r', encoding='utf-8') as f:
@@ -236,98 +165,94 @@ css="""
 .text {
   font-size: 16px;
 }
-"""
 
-# 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"
+.encryption-notice {
+  background-color: #f0f0f0;
+  padding: 15px;
+  border-radius: 5px;
+  margin-top: 10px;
+  text-align: center;
 }
+"""
 
-# Create Gradio app with enhanced security
 with gr.Blocks(css=css, elem_id="demo-container") as demo:
+    # Store encrypted data in a state variable
+    encrypted_output_state = gr.State(None)
+    
     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"
-            )
+            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"
-                    )
+                    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.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)
+                    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 Steps")
+                encrypt_password = gr.Textbox(label="Encryption Password", value="default_password", type="password")
+                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"
-            )
+            # Display placeholder image
+            image_out = gr.Image(height=800, sources=[], label="Output (Encrypted)", elem_id="output-img", format="jpg")
+            encryption_notice = gr.HTML('<div class="encryption-notice">The output image is encrypted. Use the Save button to download the encrypted file.</div>')
+            save_btn = gr.Button("Save Encrypted Image")
+            save_result = gr.Text(label="Save Result")
 
+    # Examples section
     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"]
+            ["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"))
+    
+    # Process images and encrypt outputs
+    def handle_image_generation(image, prompt, strength, seed, inference_step, encrypt_password):
+        result = process_images(image, prompt, strength, seed, inference_step, encrypt_password)
+        if result:
+            return result["display_image"], result["encrypted_data"]
+        return None, None
+    
     gr.on(
         triggers=[btn.click, prompt.submit],
-        fn=process_images,
-        inputs=[image, prompt, strength, seed, inference_step],
-        outputs=[image_out]
+        fn=handle_image_generation,
+        inputs=[image, prompt, strength, seed, inference_step, encrypt_password],
+        outputs=[image_out, encrypted_output_state]
+    )
+    
+    # Save encrypted image
+    def handle_save_encrypted(encrypted_data):
+        if encrypted_data:
+            import json
+            import tempfile
+            import os
+            
+            # Create a temporary file with the encrypted data
+            fd, path = tempfile.mkstemp(suffix='.encimg')
+            with os.fdopen(fd, 'w') as f:
+                json.dump(encrypted_data, f)
+            
+            return f"Encrypted image saved to {path}"
+        return "No encrypted image to save"
+    
+    save_btn.click(
+        fn=handle_save_encrypted,
+        inputs=[encrypted_output_state],
+        outputs=[save_result]
     )
-
-# Register shutdown handler to clean up
-import atexit
-atexit.register(secure_temp.cleanup_all)
 
 if __name__ == "__main__":
-    # Launch with security settings - FIXED: Removed 'enable_queue' parameter
-    demo.launch(
-        share=True, 
-        show_error=True,
-        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,
-        max_threads=20,
-        auth=None,             # Enable if you need authentication
-        root_path=""
-    )
+    demo.launch(share=True, show_error=True)