Upload fg-bg-klein.ipynb
Browse files
kaggle_notebooks/twin_input_setup/fg-bg-klein.ipynb
CHANGED
|
@@ -1 +1 @@
|
|
| 1 |
-
{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"datasetVersion","sourceId":16077317,"datasetId":8022630,"databundleVersionId":17046476}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =============================================================================\n# CELL 1B: Setup + Dependencies\n# =============================================================================\nfrom kaggle_secrets import UserSecretsClient\nfrom huggingface_hub import login\nimport torch\nimport os\nimport gc\nimport shutil\n\n# === Kaggle Secrets ===\nsecrets = UserSecretsClient()\ntry:\n hf_token = secrets.get_secret(\"HF_TOKEN\")\n login(token=hf_token)\n print(\"✅ HF login successful\")\nexcept Exception:\n print(\"⚠️ No HF_TOKEN secret found or login failed.\")\n\nprint(\"🧹 Removing old diffusers...\")\n!pip uninstall -y diffusers > /dev/null 2>&1\n!rm -rf /usr/local/lib/python3.12/dist-packages/diffusers* ~/.cache/pip/*diffusers*\n\nprint(\"🔄 Installing latest diffusers...\")\n!pip install -q git+https://github.com/huggingface/diffusers.git --force-reinstall --no-deps\n!python -m pip cache purge\n\nprint(\"🔐 Installing encryption + datasets + SDNQ...\")\n!pip install -q pynacl datasets sdnq\n\nprint(\"✅ Cell 1B complete! (Kaggle 2xT4 ready for 4 pipes)\")\n\n#----#\n\n# =============================================================================\n# CELL 2B: User Settings\n# =============================================================================\nkaggle_repo_path = \"/kaggle/input/datasets/nekos4lyfe/image-caption-dataset\"\npassword_txt_file = \"password.txt\"\n\n# Default values (will be overridden by txt files)\nedit_prompt = 'improve this illustration. fine art color contrast with pleasant quality. the background is dark gray.'\nresolution = '1024 x 1024 (Square)'\n\n#Hardcoded for now\ntarget_height = 1024\ntarget_width = 1024\n\nfrom pathlib import Path\nprint(\"🔑 Reading config from dataset files...\")\n\nencryption_password = Path(kaggle_repo_path, password_txt_file).read_text(encoding=\"utf-8\").strip()\nMODEL_ID = Path(kaggle_repo_path, \"model_id.txt\").read_text(encoding=\"utf-8\").strip()\nedit_prompt = Path(kaggle_repo_path, \"edit_prompt.txt\").read_text(encoding=\"utf-8\").strip()\nresolution = Path(kaggle_repo_path, \"resolution.txt\").read_text(encoding=\"utf-8\").strip()\n\nprint(\"✅ Config loaded successfully\")\nprint(f\" Model : {MODEL_ID}\")\nprint(f\" Resolution : {resolution}\")\n\nmax_image_dimension = 2048\n\n# Output settings\nsave_checkpoint_every_n = True\nsave_every_n = 30\ndebug = True\n\n# Derived paths\nencrypted_input_dataset_path = os.path.join(kaggle_repo_path, \"foregrounds\")\nbackgrounds_input_dataset_path = os.path.join(kaggle_repo_path, \"backgrounds\")\n\nif debug:\n print(f\" Foregrounds : {encrypted_input_dataset_path}\")\n print(f\" Backgrounds : {backgrounds_input_dataset_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T21:45:24.725163Z","iopub.execute_input":"2026-05-02T21:45:24.725579Z","iopub.status.idle":"2026-05-02T21:45:52.433093Z","shell.execute_reply.started":"2026-05-02T21:45:24.725552Z","shell.execute_reply":"2026-05-02T21:45:52.432020Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =============================================================================\n# CELL 3B: Stable 4 Pipes Version\n# =============================================================================\nimport torch\nimport gc\nfrom diffusers import Flux2KleinPipeline\nfrom sdnq.loader import apply_sdnq_options_to_model\nfrom functools import partial\n\ngc.collect()\ntorch.cuda.empty_cache()\n\ndef patch_encode_prompt(pipe, cached_prompt_embeds, cached_text_ids):\n def patched(self, prompt=None, device=None, num_images_per_prompt=1, prompt_embeds=None, **kwargs):\n device = device or torch.device(f\"cuda:{gpu_id}\")\n if prompt_embeds is None:\n emb = cached_prompt_embeds.to(device)\n emb = emb.repeat(1, num_images_per_prompt, 1)\n emb = emb.view(emb.shape[0] * num_images_per_prompt, emb.shape[1], -1)\n tids = cached_text_ids.to(device)\n return emb, tids\n return self._original_encode_prompt(prompt, device, num_images_per_prompt, prompt_embeds, **kwargs)\n\n if not hasattr(pipe, '_original_encode_prompt'):\n pipe._original_encode_prompt = pipe.encode_prompt\n pipe.encode_prompt = partial(patched, pipe)\n return pipe\n\n\ndef load_klein_pipe(gpu_id: int, pipe_id: int):\n torch.cuda.set_device(gpu_id)\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Loading model...\")\n \n pipe = Flux2KleinPipeline.from_pretrained(\n MODEL_ID,\n torch_dtype=torch.float16,\n low_cpu_mem_usage=True,\n device_map=\"cpu\",\n )\n\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Applying SDNQ...\")\n pipe.transformer = apply_sdnq_options_to_model(\n pipe.transformer, \n use_quantized_matmul=False # Safe on Kaggle\n )\n\n # Compute embeddings\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Computing prompt embeddings...\")\n pipe.text_encoder = pipe.text_encoder.to(f\"cuda:{gpu_id}\")\n \n with torch.inference_mode():\n prompt_embeds, text_ids = pipe.encode_prompt(\n prompt=edit_prompt,\n device=torch.device(f\"cuda:{gpu_id}\"),\n num_images_per_prompt=1,\n )\n \n pipe.edit_prompt_embeds = prompt_embeds.cpu()\n pipe.edit_text_ids = text_ids.cpu()\n pipe = patch_encode_prompt(pipe, pipe.edit_prompt_embeds, pipe.edit_text_ids)\n\n # Unload text_encoder\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Unloading text_encoder...\")\n if hasattr(pipe, \"text_encoder\"):\n pipe.text_encoder.to(\"cpu\")\n del pipe.text_encoder\n if hasattr(pipe, \"tokenizer\"):\n del pipe.tokenizer\n\n gc.collect()\n torch.cuda.empty_cache()\n\n pipe.enable_model_cpu_offload(gpu_id=gpu_id)\n pipe.vae.enable_slicing()\n pipe.vae.enable_tiling()\n pipe.enable_attention_slicing(1)\n\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Ready | VRAM: {torch.cuda.memory_allocated(gpu_id)/1e9:.2f} GB\")\n return pipe\n\n\nprint(\"🚀 Loading 4 pipes...\")\npipe00 = load_klein_pipe(0, 0)\npipe01 = load_klein_pipe(0, 1)\npipe10 = load_klein_pipe(1, 0)\npipe11 = load_klein_pipe(1, 1)\n\nprint(\"\\n✅ 4 pipes ready!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T21:45:52.434910Z","iopub.execute_input":"2026-05-02T21:45:52.435501Z","iopub.status.idle":"2026-05-02T21:47:37.610047Z","shell.execute_reply.started":"2026-05-02T21:45:52.435468Z","shell.execute_reply":"2026-05-02T21:47:37.609397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =============================================================================\n# CELL 4B: 4 Parallel Workers - ONLY ENCRYPTED OUTPUTS\n# =============================================================================\nimport os, shutil, torch, gc, io, hashlib, traceback, random as pyrandom\nfrom PIL import Image\nfrom nacl.secret import SecretBox\nfrom datasets import Dataset\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom pathlib import Path\nimport os\nimport shutil\nimport torch\nimport gc\nimport datetime\nimport io\nimport hashlib\nimport time\nimport traceback\nimport random as pyrandom\nfrom PIL import Image\nfrom nacl.secret import SecretBox\nfrom nacl.utils import random\nfrom datasets import Dataset\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\n\n# ================= ONLY ENCRYPTED OUTPUTS =================\noutput_folders = [f\"/kaggle/working/edited_gpu{i}_pipe{j}\" for i in range(2) for j in range(2)]\nfinal_output_folder = \"/kaggle/working/final_encrypted\"\nfor p in output_folders + [final_output_folder]:\n os.makedirs(p, exist_ok=True)\n\n# Password + Box\npassword = Path(kaggle_repo_path, \"password.txt\").read_text(encoding=\"utf-8\").strip()\nkey = hashlib.sha256(password.encode()).digest()\n\n# Load encrypted files\ninput_files = sorted(Path(encrypted_input_dataset_path).glob(\"*_encrypted.bin\"))\nbg_files = sorted(Path(backgrounds_input_dataset_path).glob(\"*_encrypted.bin\"))\n\nprint(f\"📊 Processing {len(input_files)} images...\")\n\npipes = [pipe00, pipe01, pipe10, pipe11]\ngpu_ids = [0, 0, 1, 1]\n\ndef worker_thread(pipe, gpu_id, start_idx, num_images, worker_id):\n\n import torch._dynamo\n from accelerate.hooks import remove_hook_from_module\n\n # ---- HARD DISABLE DYNAMO (NO CONTEXT MANAGER) ----\n torch._dynamo.reset()\n torch._dynamo.config.suppress_errors = True\n torch._dynamo.config.disable = True\n\n #os.environ[\"TORCH_COMPILE_DISABLE\"] = \"1\"\n os.environ[\"TORCHDYNAMO_DISABLE\"] = \"1\"\n \n torch.cuda.set_device(gpu_id)\n box = SecretBox(key)\n\n def decrypt(data): \n return box.decrypt(data)\n \n def encrypt(data):\n return box.encrypt(data)\n\n def bytes_to_pil(b):\n return Image.open(io.BytesIO(b)).convert(\"RGB\")\n \n def pil_to_bytes(img):\n buf = io.BytesIO()\n img.save(buf, format=\"JPEG\", quality=90)\n return buf.getvalue()\n\n for i in range(num_images):\n global_idx = start_idx + i\n if global_idx >= len(input_files):\n break\n\n try:\n gc.collect()\n torch.cuda.empty_cache()\n\n # Load input\n with open(input_files[global_idx], \"rb\") as f:\n input_enc = f.read()\n input_image = bytes_to_pil(decrypt(input_enc))\n\n # Load random background\n bg_path = pyrandom.choice(bg_files)\n with open(bg_path, \"rb\") as f:\n bg_enc = f.read()\n bg_image = bytes_to_pil(decrypt(bg_enc))\n\n # Inference\n with torch.inference_mode(), torch.no_grad():\n result = pipe(\n prompt=edit_prompt,\n image=[input_image, bg_image],\n height=target_height,\n width=target_width,\n guidance_scale=1.0,\n num_inference_steps=4,\n generator=torch.Generator(f\"cuda:{gpu_id}\").manual_seed(42),\n output_type=\"pil\",\n ).images[0]\n\n # ONLY SAVE ENCRYPTED\n encrypted_result = encrypt(pil_to_bytes(result))\n out_path = os.path.join(output_folders[worker_id], f\"edited_{global_idx:06d}.enc\")\n \n with open(out_path, \"wb\") as f:\n f.write(encrypted_result)\n\n print(f\"✅ Worker {worker_id} (GPU {gpu_id}) → {global_idx+1}/{len(input_files)}\")\n\n del result, input_image, bg_image\n\n except Exception as e:\n print(f\"❌ Error at index {global_idx}: {e}\")\n\n return True\n\n\n# Run 4 workers\nchunk_size = (len(input_files) + 3) // 4\nwith ThreadPoolExecutor(max_workers=4) as executor:\n futures = []\n for w in range(4):\n start = w * chunk_size\n num = min(chunk_size, len(input_files) - start)\n if num > 0:\n futures.append(executor.submit(worker_thread, pipes[w], gpu_ids[w], start, num, w))\n\n for future in as_completed(futures):\n future.result()\n\n# ----- Final zip (encrypted only) --------\n# -----------------------------------------\n\n#Kaggle/working/\n# - edited_gpu0_pipe0\n# - - edited0.enc\n# - - edited23.enc\n# - edited_gpu0_pipe1\n# - - edited123.enc\n# - - edited33.enc\n# - edited_gpu1_pipe0\n# - - edited78.enc\n# - - ...\n# - edited_gpu1_pipe1\n# - - edited29.enc\n# - - ...\n# - final_encrypted\n# - - final_encrypted_outputs.zip #<-- Currently this zip is empty on finishing the code\n\n# => compile the above images into /kaggle/working/final_encrypted/\n# final_encrypted_outputs.zip\n\nimport shutil\nimport os\nfrom pathlib import Path\n\n# ================== CONFIG ==================\nfinal_output_folder = \"/kaggle/working/final_encrypted\"\nfinal_zip_name = \"final_encrypted_outputs\" # without .zip\n\noutput_folders = [\n f\"/kaggle/working/edited_gpu{i}_pipe{j}\" \n for i in range(2) \n for j in range(2)\n]\n\n# Create final output folder if it doesn't exist\nos.makedirs(final_output_folder, exist_ok=True)\n\nprint(\"🔍 Collecting encrypted files from:\")\nfor folder in output_folders:\n print(f\" - {folder}\")\n\n# Collect all .enc files\ncollected_files = []\nfor folder in output_folders:\n if not os.path.exists(folder):\n print(f\"⚠️ Folder not found: {folder}\")\n continue\n \n enc_files = list(Path(folder).glob(\"*.enc\"))\n for file in enc_files:\n collected_files.append(file)\n # Copy to final folder (preserving original filename)\n shutil.copy2(file, os.path.join(final_output_folder, file.name))\n print(f\"✅ Copied: {file.name} ← {folder}\")\n\nprint(f\"\\n📦 Total encrypted files collected: {len(collected_files)}\")\n\n# Create the final zip archive\nzip_path = os.path.join(\"/kaggle/working\", final_zip_name)\nshutil.make_archive(zip_path, 'zip', final_output_folder)\n\nprint(f\"\\n🎉 Success! Final encrypted zip created:\")\nprint(f\" 📁 {zip_path}.zip\")\nprint(f\" Contains {len(collected_files)} encrypted files\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T21:47:47.511545Z","iopub.execute_input":"2026-05-02T21:47:47.512271Z","execution_failed":"2026-05-02T21:53:09.092Z"}},"outputs":[],"execution_count":null}]}
|
|
|
|
| 1 |
+
{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"datasetVersion","sourceId":16164885,"datasetId":8022630,"databundleVersionId":17140930}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =============================================================================\n# CELL 1B: Setup + Dependencies\n# =============================================================================\nfrom kaggle_secrets import UserSecretsClient\nfrom huggingface_hub import login\nimport torch\nimport os\nimport gc\nimport shutil\n\n# === Kaggle Secrets ===\nsecrets = UserSecretsClient()\ntry:\n hf_token = secrets.get_secret(\"HF_TOKEN\")\n login(token=hf_token)\n print(\"✅ HF login successful\")\nexcept Exception:\n print(\"⚠️ No HF_TOKEN secret found or login failed.\")\n\nprint(\"🧹 Removing old diffusers...\")\n!pip uninstall -y diffusers > /dev/null 2>&1\n!rm -rf /usr/local/lib/python3.12/dist-packages/diffusers* ~/.cache/pip/*diffusers*\n\nprint(\"🔄 Installing latest diffusers...\")\n!pip install -q git+https://github.com/huggingface/diffusers.git --force-reinstall --no-deps\n!python -m pip cache purge\n\nprint(\"🔐 Installing encryption + datasets + SDNQ...\")\n!pip install -q pynacl datasets sdnq\n\nprint(\"✅ Cell 1B complete! (Kaggle 2xT4 ready for 4 pipes)\")\n\n#----#\n\n# =============================================================================\n# CELL 2B: User Settings\n# =============================================================================\nkaggle_repo_path = \"/kaggle/input/datasets/nekos4lyfe/image-caption-dataset\"\npassword_txt_file = \"password.txt\"\n\n# Default values (will be overridden by txt files)\nedit_prompt = 'improve this illustration. fine art color contrast with pleasant quality. the background is dark gray.'\nresolution = '1024 x 1024 (Square)'\n\n#Hardcoded for now\ntarget_height = 1024\ntarget_width = 1024\n\nfrom pathlib import Path\nprint(\"🔑 Reading config from dataset files...\")\n\nencryption_password = Path(kaggle_repo_path, password_txt_file).read_text(encoding=\"utf-8\").strip()\nMODEL_ID = Path(kaggle_repo_path, \"model_id.txt\").read_text(encoding=\"utf-8\").strip()\nedit_prompt = Path(kaggle_repo_path, \"edit_prompt.txt\").read_text(encoding=\"utf-8\").strip()\nresolution = Path(kaggle_repo_path, \"resolution.txt\").read_text(encoding=\"utf-8\").strip()\n\nprint(\"✅ Config loaded successfully\")\nprint(f\" Model : {MODEL_ID}\")\nprint(f\" Resolution : {resolution}\")\n\nmax_image_dimension = 2048\n\n# Output settings\nsave_checkpoint_every_n = True\nsave_every_n = 30\ndebug = True\n\n# Derived paths\nencrypted_input_dataset_path = os.path.join(kaggle_repo_path, \"foregrounds\")\nbackgrounds_input_dataset_path = os.path.join(kaggle_repo_path, \"backgrounds\")\n\nif debug:\n print(f\" Foregrounds : {encrypted_input_dataset_path}\")\n print(f\" Backgrounds : {backgrounds_input_dataset_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T21:45:24.725163Z","iopub.execute_input":"2026-05-02T21:45:24.725579Z","iopub.status.idle":"2026-05-02T21:45:52.433093Z","shell.execute_reply.started":"2026-05-02T21:45:24.725552Z","shell.execute_reply":"2026-05-02T21:45:52.432020Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =============================================================================\n# CELL 3B: Stable 4 Pipes Version\n# =============================================================================\nimport torch\nimport gc\nfrom diffusers import Flux2KleinPipeline\nfrom sdnq.loader import apply_sdnq_options_to_model\nfrom functools import partial\n\ngc.collect()\ntorch.cuda.empty_cache()\n\ndef patch_encode_prompt(pipe, cached_prompt_embeds, cached_text_ids):\n def patched(self, prompt=None, device=None, num_images_per_prompt=1, prompt_embeds=None, **kwargs):\n device = device or torch.device(f\"cuda:{gpu_id}\")\n if prompt_embeds is None:\n emb = cached_prompt_embeds.to(device)\n emb = emb.repeat(1, num_images_per_prompt, 1)\n emb = emb.view(emb.shape[0] * num_images_per_prompt, emb.shape[1], -1)\n tids = cached_text_ids.to(device)\n return emb, tids\n return self._original_encode_prompt(prompt, device, num_images_per_prompt, prompt_embeds, **kwargs)\n\n if not hasattr(pipe, '_original_encode_prompt'):\n pipe._original_encode_prompt = pipe.encode_prompt\n pipe.encode_prompt = partial(patched, pipe)\n return pipe\n\n\ndef load_klein_pipe(gpu_id: int, pipe_id: int):\n torch.cuda.set_device(gpu_id)\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Loading model...\")\n \n pipe = Flux2KleinPipeline.from_pretrained(\n MODEL_ID,\n torch_dtype=torch.float16,\n low_cpu_mem_usage=True,\n device_map=\"cpu\",\n )\n\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Applying SDNQ...\")\n pipe.transformer = apply_sdnq_options_to_model(\n pipe.transformer, \n use_quantized_matmul=False # Safe on Kaggle\n )\n\n # Compute embeddings\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Computing prompt embeddings...\")\n pipe.text_encoder = pipe.text_encoder.to(f\"cuda:{gpu_id}\")\n \n with torch.inference_mode():\n prompt_embeds, text_ids = pipe.encode_prompt(\n prompt=edit_prompt,\n device=torch.device(f\"cuda:{gpu_id}\"),\n num_images_per_prompt=1,\n )\n \n pipe.edit_prompt_embeds = prompt_embeds.cpu()\n pipe.edit_text_ids = text_ids.cpu()\n pipe = patch_encode_prompt(pipe, pipe.edit_prompt_embeds, pipe.edit_text_ids)\n\n # Unload text_encoder\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Unloading text_encoder...\")\n if hasattr(pipe, \"text_encoder\"):\n pipe.text_encoder.to(\"cpu\")\n del pipe.text_encoder\n if hasattr(pipe, \"tokenizer\"):\n del pipe.tokenizer\n\n gc.collect()\n torch.cuda.empty_cache()\n\n pipe.enable_model_cpu_offload(gpu_id=gpu_id)\n pipe.vae.enable_slicing()\n pipe.vae.enable_tiling()\n pipe.enable_attention_slicing(1)\n\n print(f\" [GPU {gpu_id} - Pipe {pipe_id}] Ready | VRAM: {torch.cuda.memory_allocated(gpu_id)/1e9:.2f} GB\")\n return pipe\n\n\nprint(\"🚀 Loading 4 pipes...\")\npipe00 = load_klein_pipe(0, 0)\npipe01 = load_klein_pipe(0, 1)\npipe10 = load_klein_pipe(1, 0)\npipe11 = load_klein_pipe(1, 1)\n\nprint(\"\\n✅ 4 pipes ready!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T21:45:52.434910Z","iopub.execute_input":"2026-05-02T21:45:52.435501Z","iopub.status.idle":"2026-05-02T21:47:37.610047Z","shell.execute_reply.started":"2026-05-02T21:45:52.435468Z","shell.execute_reply":"2026-05-02T21:47:37.609397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =============================================================================\n# CELL 4B: 4 Parallel Workers - ONLY ENCRYPTED OUTPUTS\n# =============================================================================\nimport os, shutil, torch, gc, io, hashlib, traceback, random as pyrandom\nfrom PIL import Image\nfrom nacl.secret import SecretBox\nfrom datasets import Dataset\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom pathlib import Path\nimport os\nimport shutil\nimport torch\nimport gc\nimport datetime\nimport io\nimport hashlib\nimport time\nimport traceback\nimport random as pyrandom\nfrom PIL import Image\nfrom nacl.secret import SecretBox\nfrom nacl.utils import random\nfrom datasets import Dataset\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\n\n# ================= ONLY ENCRYPTED OUTPUTS =================\noutput_folders = [f\"/kaggle/working/edited_gpu{i}_pipe{j}\" for i in range(2) for j in range(2)]\nfinal_output_folder = \"/kaggle/working/final_encrypted\"\nfor p in output_folders + [final_output_folder]:\n os.makedirs(p, exist_ok=True)\n\n# Password + Box\npassword = Path(kaggle_repo_path, \"password.txt\").read_text(encoding=\"utf-8\").strip()\nkey = hashlib.sha256(password.encode()).digest()\n\n# Load encrypted files\ninput_files = sorted(Path(encrypted_input_dataset_path).glob(\"*_encrypted.bin\"))\nbg_files = sorted(Path(backgrounds_input_dataset_path).glob(\"*_encrypted.bin\"))\n\nprint(f\"📊 Processing {len(input_files)} images...\")\n\npipes = [pipe00, pipe01, pipe10, pipe11]\ngpu_ids = [0, 0, 1, 1]\n\ndef worker_thread(pipe, gpu_id, start_idx, num_images, worker_id):\n\n import torch._dynamo\n from accelerate.hooks import remove_hook_from_module\n\n # ---- HARD DISABLE DYNAMO (NO CONTEXT MANAGER) ----\n torch._dynamo.reset()\n torch._dynamo.config.suppress_errors = True\n torch._dynamo.config.disable = True\n\n #os.environ[\"TORCH_COMPILE_DISABLE\"] = \"1\"\n os.environ[\"TORCHDYNAMO_DISABLE\"] = \"1\"\n \n torch.cuda.set_device(gpu_id)\n box = SecretBox(key)\n\n def decrypt(data): \n return box.decrypt(data)\n \n def encrypt(data):\n return box.encrypt(data)\n\n def bytes_to_pil(b):\n return Image.open(io.BytesIO(b)).convert(\"RGB\")\n \n def pil_to_bytes(img):\n buf = io.BytesIO()\n img.save(buf, format=\"JPEG\", quality=90)\n return buf.getvalue()\n\n for i in range(num_images):\n global_idx = start_idx + i\n if global_idx >= len(input_files):\n break\n\n try:\n gc.collect()\n torch.cuda.empty_cache()\n\n # Load input\n with open(input_files[global_idx], \"rb\") as f:\n input_enc = f.read()\n input_image = bytes_to_pil(decrypt(input_enc))\n\n # Load random background\n bg_path = pyrandom.choice(bg_files)\n with open(bg_path, \"rb\") as f:\n bg_enc = f.read()\n bg_image = bytes_to_pil(decrypt(bg_enc))\n\n # Inference\n with torch.inference_mode(), torch.no_grad():\n result = pipe(\n prompt=edit_prompt,\n image=[input_image, bg_image],\n height=target_height,\n width=target_width,\n guidance_scale=1.0,\n num_inference_steps=4,\n generator=torch.Generator(f\"cuda:{gpu_id}\").manual_seed(42),\n output_type=\"pil\",\n ).images[0]\n\n # ONLY SAVE ENCRYPTED\n encrypted_result = encrypt(pil_to_bytes(result))\n out_path = os.path.join(output_folders[worker_id], f\"edited_{global_idx:06d}.enc\")\n \n with open(out_path, \"wb\") as f:\n f.write(encrypted_result)\n\n print(f\"✅ Worker {worker_id} (GPU {gpu_id}) → {global_idx+1}/{len(input_files)}\")\n\n del result, input_image, bg_image\n\n except Exception as e:\n print(f\"❌ Error at index {global_idx}: {e}\")\n\n return True\n\n\n# Run 4 workers\nchunk_size = (len(input_files) + 3) // 4\nwith ThreadPoolExecutor(max_workers=4) as executor:\n futures = []\n for w in range(4):\n start = w * chunk_size\n num = min(chunk_size, len(input_files) - start)\n if num > 0:\n futures.append(executor.submit(worker_thread, pipes[w], gpu_ids[w], start, num, w))\n\n for future in as_completed(futures):\n future.result()\n\n# ----- Final zip (encrypted only) --------\n# -----------------------------------------\n\n#Kaggle/working/\n# - edited_gpu0_pipe0\n# - - edited0.enc\n# - - edited23.enc\n# - edited_gpu0_pipe1\n# - - edited123.enc\n# - - edited33.enc\n# - edited_gpu1_pipe0\n# - - edited78.enc\n# - - ...\n# - edited_gpu1_pipe1\n# - - edited29.enc\n# - - ...\n# - final_encrypted\n# - - final_encrypted_outputs.zip #<-- Currently this zip is empty on finishing the code\n\n# => compile the above images into /kaggle/working/final_encrypted/\n# final_encrypted_outputs.zip\n\nimport shutil\nimport os\nfrom pathlib import Path\n\n# ================== CONFIG ==================\nfinal_output_folder = \"/kaggle/working/final_encrypted\"\nfinal_zip_name = \"final_encrypted_outputs\" # without .zip\n\noutput_folders = [\n f\"/kaggle/working/edited_gpu{i}_pipe{j}\" \n for i in range(2) \n for j in range(2)\n]\n\n# Create final output folder if it doesn't exist\nos.makedirs(final_output_folder, exist_ok=True)\n\nprint(\"🔍 Collecting encrypted files from:\")\nfor folder in output_folders:\n print(f\" - {folder}\")\n\n# Collect all .enc files\ncollected_files = []\nfor folder in output_folders:\n if not os.path.exists(folder):\n print(f\"⚠️ Folder not found: {folder}\")\n continue\n \n enc_files = list(Path(folder).glob(\"*.enc\"))\n for file in enc_files:\n collected_files.append(file)\n # Copy to final folder (preserving original filename)\n shutil.copy2(file, os.path.join(final_output_folder, file.name))\n print(f\"✅ Copied: {file.name} ← {folder}\")\n\nprint(f\"\\n📦 Total encrypted files collected: {len(collected_files)}\")\n\n# Create the final zip archive\nzip_path = os.path.join(\"/kaggle/working\", final_zip_name)\nshutil.make_archive(zip_path, 'zip', final_output_folder)\n\nprint(f\"\\n🎉 Success! Final encrypted zip created:\")\nprint(f\" 📁 {zip_path}.zip\")\nprint(f\" Contains {len(collected_files)} encrypted files\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-02T21:47:47.511545Z","iopub.execute_input":"2026-05-02T21:47:47.512271Z","execution_failed":"2026-05-02T21:53:09.092Z"}},"outputs":[],"execution_count":null}]}
|