Update app.py

app.py

@@ -1,22 +1,28 @@
-import spaces
-import gradio as gr
-import re
-from PIL import Image
-import io
-import base64
 import os
+import io
 import json
+import base64
+import re
+from PIL import Image
 import numpy as np
 import torch
+import gradio as gr
+import spaces
 from diffusers import FluxImg2ImgPipeline
 from cryptography.fernet import Fernet
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
+# Retrieve the encryption key from the environment (set in Hugging Face Secrets Manager)
+ENCRYPTION_KEY = os.environ.get("key", "FAKEFALLBACKKEY_FOR_LOCAL_TESTING")
+
 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)
+pipe = FluxImg2ImgPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", 
+    torch_dtype=torch.bfloat16
+).to(device)
 
 def generate_key(password, salt=None):
     if salt is None:
@@ -30,16 +36,19 @@ def generate_key(password, salt=None):
     key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
     return key, salt
 
-def encrypt_image(image):
+def encrypt_image(image, password=None):
+    # Use the secure key if no override is provided
+    if password is None:
+        password = ENCRYPTION_KEY
+
     # Convert PIL Image to bytes
     img_byte_arr = io.BytesIO()
     image.save(img_byte_arr, format='PNG')
     img_byte_arr = img_byte_arr.getvalue()
 
-    # Generate key for encryption using the secret from the environment
-    key, salt = generate_key(ENCRYPTION_KEY)
+    # Generate key for encryption using the secure password
+    key, salt = generate_key(password)
     cipher = Fernet(key)
-
     encrypted_data = cipher.encrypt(img_byte_arr)
 
     return {
@@ -49,15 +58,17 @@ def encrypt_image(image):
         'original_height': image.height
     }
 
-def decrypt_image(encrypted_data_dict):
+def decrypt_image(encrypted_data_dict, password=None):
+    if password is None:
+        password = ENCRYPTION_KEY
+
     # 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 same password and salt
-    key, _ = generate_key(ENCRYPTION_KEY, salt)
+    # Regenerate the key using the secure password and salt
+    key, _ = generate_key(password, salt)
     cipher = Fernet(key)
-
     decrypted_data = cipher.decrypt(encrypted_data)
     image = Image.open(io.BytesIO(decrypted_data))
     return image
@@ -72,7 +83,7 @@ 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
-    
+
     if width > height:
         scaling_factor = maximum_size / width
     else:
@@ -88,19 +99,18 @@ def adjust_to_multiple_of_32(width: int, height: int):
     return width, height
 
 @spaces.GPU(duration=120)
-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)):
+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")
 
     def process_img2img(image, prompt="a person", strength=0.75, seed=0, num_inference_steps=4):
         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)
-    
+
         output = pipe(
             prompt=prompt,
             image=image,
@@ -112,28 +122,26 @@ def process_images(image, prompt="a girl", strength=0.75, seed=0, inference_step
             num_inference_steps=num_inference_steps,
             max_sequence_length=256
         )
-    
+
         pil_image = output.images[0]
         new_width, new_height = pil_image.size
-    
+
         if (new_width != fit_width) or (new_height != fit_height):
             resized_image = pil_image.resize((fit_width, fit_height), Image.LANCZOS)
             return resized_image
         return pil_image
-    
+
     output = process_img2img(image, prompt, strength, seed, inference_step)
-    
-    # Encrypt the output image
+
+    # Encrypt the output image using the secure key
     if output is not None:
-        encrypted_output = encrypt_image(output, encrypt_password)
-        # Instead of returning a gray placeholder, show the real pipeline result:
+        encrypted_output = encrypt_image(output)
         return {
-            "display_image": output,  # <--- Use the real pipeline output
+            "display_image": output,
             "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)
@@ -157,7 +165,7 @@ css = """
   display: flex;
   align-items: center;
   justify-content: center;
-  gap:10px
+  gap:10px;
 }
 
 .image {
@@ -182,7 +190,7 @@ css = """
 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"))
@@ -204,9 +212,7 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
                         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(
@@ -218,15 +224,8 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
                     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():
-            # Display placeholder image
             image_out = gr.Image(
                 height=800,
                 sources=[],
@@ -252,31 +251,29 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
         ],
         inputs=[image, image_out, prompt],
     )
-    
+
     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)
+    def handle_image_generation(image, prompt, strength, seed, inference_step):
+        result = process_images(image, prompt, strength, seed, inference_step)
         if result:
             return result["display_image"], result["encrypted_data"]
         return None, None
 
-    # >>>> CHANGED: Use .click() and .submit() with api_name
     btn.click(
         fn=handle_image_generation,
-        inputs=[image, prompt, strength, seed, inference_step, encrypt_password],
+        inputs=[image, prompt, strength, seed, inference_step],
         outputs=[image_out, encrypted_output_state],
-        api_name="/process_images"  # Exposes handle_image_generation as /process_images
+        api_name="/process_images"
     )
 
     prompt.submit(
         fn=handle_image_generation,
-        inputs=[image, prompt, strength, seed, inference_step, encrypt_password],
+        inputs=[image, prompt, strength, seed, inference_step],
         outputs=[image_out, encrypted_output_state],
-        api_name="/process_images"  # Same endpoint
+        api_name="/process_images"
     )
-    # <<<< END CHANGE
 
     def handle_save_encrypted(encrypted_data):
         if encrypted_data:
@@ -286,7 +283,7 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
                 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],