Upload 2 files

app.py

@@ -0,0 +1,316 @@
+import io
+import os
+import time
+import mimetypes
+
+import gradio as gr
+import requests
+import torch
+from diffusers import AutoPipelineForText2Image
+from PIL import Image
+
+# --- Configuration ---
+# Fetch the API key from environment variables.
+OMNISEAL_API_KEY = os.getenv("OMNISEAL_API_KEY")
+if not OMNISEAL_API_KEY:
+    raise ValueError("Error: OMNISEAL_API_KEY environment variable not set.")
+
+# Updated API endpoints to use the unified direct API.
+OMNISEAL_EMBED_URL = "https://api.purecipher.com/unified/embed-file"
+OMNISEAL_RECOVER_URL = "https://api.purecipher.com/unified/retrieve-file"
+
+# --- Model Loading ---
+# Load the model once when the script starts.
+print("Loading image generation model... This might take a few minutes the first time.")
+try:
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        "stabilityai/sd-turbo", torch_dtype=torch.float32
+    )
+    pipe.to("cpu")
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    pipe = None
+
+# --- Helper Functions ---
+def generate_image_from_prompt(prompt: str) -> Image:
+    """
+    Generates an image locally using the stabilityai/sd-turbo model.
+    Returns a PIL Image.
+    """
+    if pipe is None:
+        raise gr.Error(
+            "Image generation model is not loaded. Please check the console for errors."
+        )
+    try:
+        image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
+        return image
+    except Exception as e:
+        raise gr.Error(f"Error generating image locally: {e}")
+
+def call_omni_seal_api(url: str, files: dict, data: dict) -> bytes:
+    """
+    Makes a POST request to the OmniSeal API, automatically including the API token.
+    Returns the response content as bytes.
+    """
+    headers = {
+        "accept": "application/json",
+        "api_token": OMNISEAL_API_KEY
+    }
+    try:
+        response = requests.post(url, files=files, data=data, headers=headers, timeout=120)
+        response.raise_for_status()
+        return response.content
+    except requests.exceptions.RequestException as e:
+        error_message = f"API call failed: {e}."
+        if e.response is not None:
+             error_message += f" Response: {e.response.text}"
+        raise gr.Error(error_message)
+
+def pil_to_bytes(image: Image, format="PNG") -> bytes:
+    """Converts a PIL Image to bytes."""
+    with io.BytesIO() as buf:
+        image.save(buf, format=format)
+        return buf.getvalue()
+
+def bytes_to_pil(image_bytes: bytes) -> Image:
+    """Converts bytes to a PIL Image."""
+    try:
+        return Image.open(io.BytesIO(image_bytes))
+    except Exception:
+        # If the returned file isn't an image, we can't display it, which is expected for recovery.
+        return None
+
+# --- Gradio Handler Functions ---
+def fn_generate_image(prompt):
+    gr.Info("Generating image... This may take a moment.")
+    image = generate_image_from_prompt(prompt)
+    gr.Info("Image generated successfully!")
+    return image, image
+
+def fn_embed_secret(generated_image_state, secret_file, private_seed):
+    if not private_seed:
+        raise gr.Error("Private seed is required.")
+    if generated_image_state is None:
+        raise gr.Error("Please generate an image first.")
+    if secret_file is None:
+        raise gr.Error("Please upload a secret file.")
+
+    gr.Info("Calling OmniSeal API to embed secret...")
+    generated_image_bytes = pil_to_bytes(generated_image_state)
+    
+    secret_file_name = os.path.basename(secret_file.name)
+    secret_mime_type, _ = mimetypes.guess_type(secret_file.name)
+    if secret_mime_type is None:
+        secret_mime_type = "application/octet-stream"
+
+    with open(secret_file.name, "rb") as f:
+        secret_file_bytes = f.read()
+
+    files = {
+        "cover_file": ("generated_image.png", generated_image_bytes, "image/png"),
+        "secret_file": (secret_file_name, secret_file_bytes, secret_mime_type),
+    }
+    data = {"private_seed": private_seed}
+
+    embedded_image_bytes = call_omni_seal_api(OMNISEAL_EMBED_URL, files, data)
+    embedded_image_pil = bytes_to_pil(embedded_image_bytes)
+    
+    embedded_filename = f"embedded_{secret_file_name}.png"
+    with open(embedded_filename, "wb") as f:
+        f.write(embedded_image_bytes)
+
+    gr.Info("Secret file embedded successfully!")
+    return embedded_image_pil, embedded_filename, embedded_image_pil, secret_file_name
+
+def fn_recover_secret(embedded_filename, secret_file_name_state, private_seed):
+    if not private_seed:
+        raise gr.Error("Private seed is required.")
+    if embedded_filename is None:
+        raise gr.Error("No embedded image available to recover from.")
+
+    gr.Info("Calling OmniSeal API to recover secret...")
+    #embedded_image_bytes = pil_to_bytes(embedded_image_state)
+    uploaded_file_name = os.path.basename(embedded_filename.name)
+    uploaded_mime_type, _ = mimetypes.guess_type(embedded_filename.name)
+    if uploaded_mime_type is None:
+        uploaded_mime_type = "application/octet-stream"
+
+    with open(embedded_filename.name, "rb") as f:
+        embedded_image_bytes = f.read()
+    files = {"sealed_file": ("embedded_image.png", embedded_image_bytes, "image/png")}
+    data = {"private_seed": private_seed}
+
+    recovered_bytes = call_omni_seal_api(OMNISEAL_RECOVER_URL, files, data)
+    
+    recovered_filename = secret_file_name_state if secret_file_name_state else "recovered_secret.bin"
+    with open(recovered_filename, "wb") as f:
+        f.write(recovered_bytes)
+
+    gr.Info("Secret recovered successfully!")
+    return recovered_filename
+
+def fn_recover_from_upload(uploaded_file, private_seed):
+    if not private_seed:
+        raise gr.Error("Private seed is required.")
+    if uploaded_file is None:
+        raise gr.Error("Please upload an embedded image file.")
+
+    gr.Info("Calling OmniSeal API to recover secret from uploaded file...")
+    uploaded_file_name = os.path.basename(uploaded_file.name)
+    uploaded_mime_type, _ = mimetypes.guess_type(uploaded_file.name)
+    if uploaded_mime_type is None:
+        uploaded_mime_type = "application/octet-stream"
+
+    with open(uploaded_file.name, "rb") as f:
+        embedded_bytes = f.read()
+        
+    files = {"sealed_file": (uploaded_file_name, embedded_bytes, uploaded_mime_type)}
+    data = {"private_seed": private_seed}
+    
+    recovered_bytes = call_omni_seal_api(OMNISEAL_RECOVER_URL, files, data)
+    
+    recovered_filename = f"recovered_from_{uploaded_file_name}"
+    recovered_filename = ".".join(recovered_filename.split(".")[:-1]) if "." in recovered_filename else recovered_filename
+
+    with open(recovered_filename, "wb") as f:
+        f.write(recovered_bytes)
+        
+    gr.Info("Secret recovered successfully!")
+    return recovered_filename
+    
+def fn_embed_from_upload(cover_file, secret_file, private_seed):
+    if not private_seed:
+        raise gr.Error("Private seed is required.")
+    if cover_file is None:
+        raise gr.Error("Please upload a cover file.")
+    if secret_file is None:
+        raise gr.Error("Please upload a secret file.")
+
+    gr.Info("Calling OmniSeal API to embed secret in uploaded file...")
+    
+    cover_file_name = os.path.basename(cover_file.name)
+    cover_mime_type, _ = mimetypes.guess_type(cover_file.name)
+    if cover_mime_type is None:
+        cover_mime_type = "application/octet-stream"
+        
+    secret_file_name = os.path.basename(secret_file.name)
+    secret_mime_type, _ = mimetypes.guess_type(secret_file.name)
+    if secret_mime_type is None:
+        secret_mime_type = "application/octet-stream"
+
+    with open(cover_file.name, "rb") as f_cover, open(secret_file.name, "rb") as f_secret:
+        cover_file_bytes = f_cover.read()
+        secret_file_bytes = f_secret.read()
+
+    files = {
+        "cover_file": (cover_file_name, cover_file_bytes, cover_mime_type),
+        "secret_file": (secret_file_name, secret_file_bytes, secret_mime_type),
+    }
+    data = {"private_seed": private_seed}
+
+    embedded_file_bytes = call_omni_seal_api(OMNISEAL_EMBED_URL, files, data)
+    embedded_file_pil = bytes_to_pil(embedded_file_bytes)
+    
+    cover_ext = os.path.splitext(cover_file_name)[1]
+    embedded_filename = f"embedded_{secret_file_name}{cover_ext}"
+    with open(embedded_filename, "wb") as f:
+        f.write(embedded_file_bytes)
+
+    gr.Info("Secret file embedded successfully!")
+    return embedded_file_pil, embedded_filename
+
+# --- Gradio UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Image(value="purecipher_logo_white.png", label="PureCipher Logo", show_label=False)
+    gr.Markdown("# PureCipher OmniSeal Demo")
+    
+    gr.Markdown(
+        """
+    This application demonstrates the power of **PureCipher's OmniSeal API** to hide a secret file within another file.
+    1.  **Generate an Image**: Type a prompt to create an image using AI.
+    2.  **Embed a Secret**: Upload any file to hide it inside the generated image.
+    3.  **Recover the Secret**: Extract the original secret file from the new image.
+    """
+    )
+    
+    with gr.Row():
+        private_seed_input = gr.Textbox(label="Private Seed", value="1234", type="password")
+
+    # State variables to hold data between steps
+    generated_image_state = gr.State()
+    embedded_image_state = gr.State()
+    secret_file_name_state = gr.State()
+
+    with gr.Tabs():
+        with gr.TabItem("Generate & Embed"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("### Step 1: Generate Image")
+                    prompt_input = gr.Textbox(label="Prompt", value="A majestic lion in the savanna")
+                    generate_btn = gr.Button("Generate Image", variant="primary")
+                    gr.Markdown("---")
+                    gr.Markdown("### Step 2: Embed Secret File")
+                    secret_file_input = gr.File(label="Upload Secret File")
+                    embed_btn = gr.Button("Embed Secret File", variant="primary")
+                with gr.Column(scale=2):
+                    gr.Markdown("### Results")
+                    with gr.Row():
+                        generated_image_output = gr.Image(label="Original Generated Image", type="pil")
+                        embedded_image_output = gr.Image(label="Image with Embedded Secret", type="pil")
+                    recover_btn = gr.Button("Recover Secret from this Image")
+                    with gr.Row():
+                        embedded_file_output = gr.File(label="Download Embedded Image")
+                        recovered_file_output = gr.File(label="Download Recovered Secret")
+
+        with gr.TabItem("Embed & Recover"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("### Step 1: Upload Files")
+                    cover_file_input_upload = gr.File(label="Upload Cover File")
+                    secret_file_input_upload = gr.File(label="Upload Secret File")
+                    embed_upload_btn = gr.Button("Embed Secret File", variant="primary")
+                with gr.Column(scale=2):
+                    gr.Markdown("### Results")
+                    embedded_image_output_upload = gr.Image(label="Image with Embedded Secret", type="pil")
+                    embedded_file_output_upload = gr.File(label="Download Embedded File")
+
+        with gr.TabItem("Recover from Existing Image"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Upload an image that already has a secret file embedded in it.")
+                    upload_image_input = gr.File(label="Upload Embedded Image")
+                    recover_upload_btn = gr.Button("Recover Secret from Uploaded Image", variant="primary")
+                with gr.Column():
+                    gr.Markdown("### Recovered File")
+                    upload_recovered_output = gr.File(label="Download Recovered Secret")
+
+    # Define interactions
+    generate_btn.click(
+        fn=fn_generate_image,
+        inputs=[prompt_input],
+        outputs=[generated_image_output, generated_image_state],
+    )
+    embed_btn.click(
+        fn=fn_embed_secret,
+        inputs=[generated_image_state, secret_file_input, private_seed_input],
+        outputs=[embedded_image_output, embedded_file_output, embedded_image_state, secret_file_name_state],
+    )
+    recover_btn.click(
+        fn=fn_recover_secret,
+        inputs=[embedded_file_output, secret_file_name_state, private_seed_input],
+        outputs=[recovered_file_output],
+    )
+    recover_upload_btn.click(
+        fn=fn_recover_from_upload,
+        inputs=[upload_image_input, private_seed_input],
+        outputs=[upload_recovered_output],
+    )
+    embed_upload_btn.click(
+        fn=fn_embed_from_upload,
+        inputs=[cover_file_input_upload, secret_file_input_upload, private_seed_input],
+        outputs=[embedded_image_output_upload, embedded_file_output_upload],
+    )
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=["purecipher_logo_white.png"])

requirements.txt

@@ -0,0 +1,7 @@
+accelerate
+diffusers
+invisible_watermark
+torch
+transformers
+xformers
+pillow==11.2.1