{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "gpuType": "T4" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "PwuBLMVifFoV" }, "outputs": [], "source": [ "#@markdown =============================================================================\n", "#@markdown CELL 1: Google Colab Setup + Dependencies\n", "#@markdown =============================================================================\n", "\n", "# --- Mount Google Drive ---\n", "from google.colab import drive\n", "drive.mount('/content/drive')\n", "\n", "# --- HF Login ---\n", "from google.colab import userdata\n", "from huggingface_hub import login\n", "\n", "try:\n", " hf_token = userdata.get('HF_TOKEN')\n", " login(token=hf_token)\n", " print(\"โœ… HF login successful\")\n", "except Exception as e:\n", " print(f\"โš ๏ธ HF login failed: {e}\")\n", "\n", "# --- Cleanup old diffusers ---\n", "import os\n", "import gc\n", "import shutil\n", "import torch\n", "\n", "print(\"๐Ÿงน Removing old diffusers...\")\n", "!pip uninstall -y diffusers > /dev/null 2>&1\n", "!rm -rf /usr/local/lib/python3.11/dist-packages/diffusers*\n", "!rm -rf ~/.cache/pip/*diffusers*\n", "\n", "print(\"๐Ÿ”„ Installing latest diffusers...\")\n", "!pip install -q git+https://github.com/huggingface/diffusers.git --force-reinstall --no-deps\n", "\n", "print(\"๐Ÿ” Installing dependencies...\")\n", "!pip install -q pynacl datasets accelerate transformers sentencepiece safetensors sdnq\n", "\n", "print(\"๐Ÿงน Purging pip cache...\")\n", "!python -m pip cache purge\n", "\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "print(\"โœ… Cell 1 complete!\")" ] }, { "cell_type": "code", "source": [ "#@markdown =============================================================================\n", "#@markdown CELL 2: User Settings + Google Drive Dataset Paths\n", "#@markdown =============================================================================\n", "\n", "import os\n", "from pathlib import Path\n", "\n", "# =============================================================================\n", "# GOOGLE DRIVE DATASET ROOT\n", "# =============================================================================\n", "\n", "# Example:\n", "# /content/drive/MyDrive/my_flux_dataset/\n", "#\n", "# โ”œโ”€โ”€ password.txt\n", "# โ”œโ”€โ”€ model_id.txt\n", "# โ”œโ”€โ”€ edit_prompt.txt\n", "# โ”œโ”€โ”€ resolution.txt\n", "# โ”œโ”€โ”€ foregrounds/\n", "# โ””โ”€โ”€ backgrounds/\n", "\n", "google_drive_dataset_path = \"/content/drive/MyDrive/my_flux_dataset\"\n", "\n", "# =============================================================================\n", "# CONFIG FILES\n", "# =============================================================================\n", "\n", "password_txt_file = \"password.txt\"\n", "\n", "# Defaults\n", "edit_prompt = \"improve this illustration\"\n", "resolution = \"1024 x 1024 (Square)\"\n", "\n", "# Hardcoded dimensions\n", "target_height = 1024\n", "target_width = 1024\n", "\n", "print(\"๐Ÿ”‘ Reading config files from Google Drive...\")\n", "\n", "dataset_path = Path(google_drive_dataset_path)\n", "\n", "encryption_password = (dataset_path / \"password.txt\").read_text(encoding=\"utf-8\").strip()\n", "MODEL_ID = (dataset_path / \"model_id.txt\").read_text(encoding=\"utf-8\").strip()\n", "edit_prompt = (dataset_path / \"edit_prompt.txt\").read_text(encoding=\"utf-8\").strip()\n", "resolution = (dataset_path / \"resolution.txt\").read_text(encoding=\"utf-8\").strip()\n", "\n", "print(\"โœ… Config loaded successfully\")\n", "print(f\" Model : {MODEL_ID}\")\n", "print(f\" Resolution : {resolution}\")\n", "\n", "# =============================================================================\n", "# DATASET PATHS\n", "# =============================================================================\n", "\n", "encrypted_input_dataset_path = os.path.join(\n", " google_drive_dataset_path,\n", " \"foregrounds.zip\"\n", ")\n", "\n", "backgrounds_input_dataset_path = os.path.join(\n", " google_drive_dataset_path,\n", " \"backgrounds.zip\"\n", ")\n", "\n", "print(f\" Foregrounds : {encrypted_input_dataset_path}\")\n", "print(f\" Backgrounds : {backgrounds_input_dataset_path}\")\n", "\n", "# =============================================================================\n", "# SETTINGS\n", "# =============================================================================\n", "\n", "debug = True\n", "save_checkpoint_every_n = False\n", "\n", "# =============================================================================\n", "# MULTI GPU SUPPORT\n", "# =============================================================================\n", "\n", "NUM_GPUS = torch.cuda.device_count()\n", "\n", "if NUM_GPUS == 0:\n", " raise RuntimeError(\"โŒ No CUDA GPU detected\")\n", "\n", "print(f\"\\n๐Ÿš€ Detected {NUM_GPUS} GPU(s)\")\n", "\n", "for i in range(NUM_GPUS):\n", " print(f\" GPU {i}: {torch.cuda.get_device_name(i)}\")\n", "\n", "# Pipes per GPU\n", "PIPES_PER_GPU = 2\n", "\n", "TOTAL_PIPES = NUM_GPUS * PIPES_PER_GPU\n", "\n", "print(f\"\\n๐Ÿง  Total Pipes Configured: {TOTAL_PIPES}\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "qhp3jsDwfGom", "outputId": "4b4882ca-5f6d-435f-b21d-e7c9644a1a82" }, "execution_count": 2, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "๐Ÿ”‘ Reading config files from Google Drive...\n", "โœ… Config loaded successfully\n", " Model : codeShare/FLUX.2-klein-AIO-SDNQ-4bit-dynamic\n", " Resolution : 1024 x 1024 (Square)\n", " Foregrounds : /content/drive/MyDrive/my_flux_dataset/foregrounds\n", " Backgrounds : /content/drive/MyDrive/my_flux_dataset/backgrounds\n", "\n", "๐Ÿš€ Detected 1 GPU(s)\n", " GPU 0: Tesla T4\n", "\n", "๐Ÿง  Total Pipes Configured: 2\n" ] } ] }, { "cell_type": "code", "source": [ "#@markdown =============================================================================\n", "#@markdown CELL 3: Load Flux Klein Pipes\n", "#@markdown =============================================================================\n", "\n", "import gc\n", "import torch\n", "from functools import partial\n", "\n", "from diffusers import Flux2KleinPipeline\n", "from sdnq.loader import apply_sdnq_options_to_model\n", "\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "# =============================================================================\n", "# PATCH PROMPT ENCODING\n", "# =============================================================================\n", "\n", "def patch_encode_prompt(pipe, cached_prompt_embeds, cached_text_ids, gpu_id):\n", "\n", " def patched(\n", " self,\n", " prompt=None,\n", " device=None,\n", " num_images_per_prompt=1,\n", " prompt_embeds=None,\n", " **kwargs\n", " ):\n", " device = device or torch.device(f\"cuda:{gpu_id}\")\n", "\n", " if prompt_embeds is None:\n", " emb = cached_prompt_embeds.to(device)\n", "\n", " emb = emb.repeat(1, num_images_per_prompt, 1)\n", "\n", " emb = emb.view(\n", " emb.shape[0] * num_images_per_prompt,\n", " emb.shape[1],\n", " -1\n", " )\n", "\n", " tids = cached_text_ids.to(device)\n", "\n", " return emb, tids\n", "\n", " return self._original_encode_prompt(\n", " prompt,\n", " device,\n", " num_images_per_prompt,\n", " prompt_embeds,\n", " **kwargs\n", " )\n", "\n", " if not hasattr(pipe, \"_original_encode_prompt\"):\n", " pipe._original_encode_prompt = pipe.encode_prompt\n", "\n", " pipe.encode_prompt = partial(patched, pipe)\n", "\n", " return pipe\n", "\n", "# =============================================================================\n", "# LOAD PIPE\n", "# =============================================================================\n", "\n", "def load_klein_pipe(gpu_id: int, pipe_id: int):\n", "\n", " torch.cuda.set_device(gpu_id)\n", "\n", " print(f\"\\n๐Ÿš€ Loading Pipe {pipe_id} on GPU {gpu_id}\")\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\"๐Ÿ”ง Applying SDNQ to Pipe {pipe_id}\")\n", "\n", " pipe.transformer = apply_sdnq_options_to_model(\n", " pipe.transformer,\n", " use_quantized_matmul=False\n", " )\n", "\n", " # -------------------------------------------------------------------------\n", " # Prompt Embeddings\n", " # -------------------------------------------------------------------------\n", "\n", " print(f\"๐Ÿง  Computing prompt embeddings for Pipe {pipe_id}\")\n", "\n", " pipe.text_encoder = pipe.text_encoder.to(f\"cuda:{gpu_id}\")\n", "\n", " with torch.inference_mode():\n", "\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", "\n", " pipe = patch_encode_prompt(\n", " pipe,\n", " pipe.edit_prompt_embeds,\n", " pipe.edit_text_ids,\n", " gpu_id\n", " )\n", "\n", " # -------------------------------------------------------------------------\n", " # Remove text encoder after caching embeddings\n", " # -------------------------------------------------------------------------\n", "\n", " print(f\"๐Ÿงน Unloading text encoder for Pipe {pipe_id}\")\n", "\n", " if hasattr(pipe, \"text_encoder\"):\n", " pipe.text_encoder.to(\"cpu\")\n", " del pipe.text_encoder\n", "\n", " if hasattr(pipe, \"tokenizer\"):\n", " del pipe.tokenizer\n", "\n", " gc.collect()\n", " torch.cuda.empty_cache()\n", "\n", " # -------------------------------------------------------------------------\n", " # CPU Offload\n", " # -------------------------------------------------------------------------\n", "\n", " pipe.enable_model_cpu_offload(gpu_id=gpu_id)\n", "\n", " pipe.vae.enable_slicing()\n", " pipe.vae.enable_tiling()\n", "\n", " pipe.enable_attention_slicing(1)\n", "\n", " print(\n", " f\"โœ… Pipe {pipe_id} Ready | \"\n", " f\"GPU {gpu_id} VRAM: \"\n", " f\"{torch.cuda.memory_allocated(gpu_id)/1e9:.2f} GB\"\n", " )\n", "\n", " return pipe\n", "\n", "# =============================================================================\n", "# CREATE PIPES\n", "# =============================================================================\n", "\n", "pipes = []\n", "gpu_ids = []\n", "\n", "pipe_counter = 0\n", "\n", "for gpu_id in range(NUM_GPUS):\n", "\n", " for local_pipe_id in range(PIPES_PER_GPU):\n", "\n", " pipe = load_klein_pipe(gpu_id, pipe_counter)\n", "\n", " pipes.append(pipe)\n", " gpu_ids.append(gpu_id)\n", "\n", " pipe_counter += 1\n", "\n", "print(\"\\n๐ŸŽ‰ ALL PIPES READY!\")" ], "metadata": { "id": "ZlFpC5tmfJX8" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown =============================================================================\n", "#@markdown CELL 4: Parallel Image Editing Workers (ZIP Dataset Version)\n", "#@markdown =============================================================================\n", "\n", "import os\n", "import io\n", "import gc\n", "import zipfile\n", "import torch\n", "import shutil\n", "import hashlib\n", "import traceback\n", "import random as pyrandom\n", "\n", "from pathlib import Path\n", "from PIL import Image\n", "from concurrent.futures import ThreadPoolExecutor, as_completed\n", "\n", "from nacl.secret import SecretBox\n", "\n", "# =============================================================================\n", "# OUTPUT DIRECTORIES\n", "# =============================================================================\n", "\n", "output_folders = [\n", " f\"/content/edited_gpu{gpu_ids[i]}_pipe{i}\"\n", " for i in range(len(pipes))\n", "]\n", "\n", "final_output_folder = \"/content/final_encrypted\"\n", "\n", "for folder in output_folders + [final_output_folder]:\n", " os.makedirs(folder, exist_ok=True)\n", "\n", "# =============================================================================\n", "# ENCRYPTION SETUP\n", "# =============================================================================\n", "\n", "password = Path(\n", " google_drive_dataset_path,\n", " \"password.txt\"\n", ").read_text(encoding=\"utf-8\").strip()\n", "\n", "key = hashlib.sha256(password.encode()).digest()\n", "\n", "# =============================================================================\n", "# OPEN ZIP FILES\n", "# =============================================================================\n", "\n", "print(\"๐Ÿ“ฆ Opening ZIP datasets...\")\n", "\n", "foreground_zip = zipfile.ZipFile(\n", " encrypted_input_dataset_path,\n", " \"r\"\n", ")\n", "\n", "background_zip = zipfile.ZipFile(\n", " backgrounds_input_dataset_path,\n", " \"r\"\n", ")\n", "\n", "# =============================================================================\n", "# LOAD FILE LISTS\n", "# =============================================================================\n", "\n", "input_files = sorted([\n", " x for x in foreground_zip.namelist()\n", " if x.endswith(\"_encrypted.bin\")\n", "])\n", "\n", "bg_files = sorted([\n", " x for x in background_zip.namelist()\n", " if x.endswith(\"_encrypted.bin\")\n", "])\n", "\n", "print(f\"๐Ÿ“Š Found {len(input_files)} foreground images\")\n", "print(f\"๐ŸŒ„ Found {len(bg_files)} background images\")\n", "\n", "if len(input_files) == 0:\n", " raise RuntimeError(\"โŒ No foreground encrypted files found\")\n", "\n", "if len(bg_files) == 0:\n", " raise RuntimeError(\"โŒ No background encrypted files found\")\n", "\n", "# =============================================================================\n", "# WORKER FUNCTION\n", "# =============================================================================\n", "\n", "def worker_thread(\n", " pipe,\n", " gpu_id,\n", " start_idx,\n", " num_images,\n", " worker_id\n", "):\n", "\n", " import torch._dynamo\n", "\n", " torch._dynamo.reset()\n", " torch._dynamo.config.suppress_errors = True\n", " torch._dynamo.config.disable = True\n", "\n", " os.environ[\"TORCHDYNAMO_DISABLE\"] = \"1\"\n", "\n", " torch.cuda.set_device(gpu_id)\n", "\n", " box = SecretBox(key)\n", "\n", " # -------------------------------------------------------------------------\n", " # Reopen ZIPs inside thread\n", " # -------------------------------------------------------------------------\n", "\n", " local_foreground_zip = zipfile.ZipFile(\n", " encrypted_input_dataset_path,\n", " \"r\"\n", " )\n", "\n", " local_background_zip = zipfile.ZipFile(\n", " backgrounds_input_dataset_path,\n", " \"r\"\n", " )\n", "\n", " # -------------------------------------------------------------------------\n", " # Encryption Helpers\n", " # -------------------------------------------------------------------------\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", "\n", " buf = io.BytesIO()\n", "\n", " img.save(\n", " buf,\n", " format=\"JPEG\",\n", " quality=90\n", " )\n", "\n", " return buf.getvalue()\n", "\n", " # -------------------------------------------------------------------------\n", " # Processing Loop\n", " # -------------------------------------------------------------------------\n", "\n", " for i in range(num_images):\n", "\n", " global_idx = start_idx + i\n", "\n", " if global_idx >= len(input_files):\n", " break\n", "\n", " try:\n", "\n", " gc.collect()\n", " torch.cuda.empty_cache()\n", "\n", " # -------------------------------------------------------------\n", " # Load encrypted foreground from ZIP\n", " # -------------------------------------------------------------\n", "\n", " fg_name = input_files[global_idx]\n", "\n", " with local_foreground_zip.open(fg_name) as f:\n", " input_enc = f.read()\n", "\n", " input_image = bytes_to_pil(\n", " decrypt(input_enc)\n", " )\n", "\n", " # -------------------------------------------------------------\n", " # Load random encrypted background from ZIP\n", " # -------------------------------------------------------------\n", "\n", " bg_name = pyrandom.choice(bg_files)\n", "\n", " with local_background_zip.open(bg_name) as f:\n", " bg_enc = f.read()\n", "\n", " bg_image = bytes_to_pil(\n", " decrypt(bg_enc)\n", " )\n", "\n", " # -------------------------------------------------------------\n", " # Inference\n", " # -------------------------------------------------------------\n", "\n", " with torch.inference_mode(), torch.no_grad():\n", "\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(\n", " f\"cuda:{gpu_id}\"\n", " ).manual_seed(42),\n", " output_type=\"pil\",\n", " ).images[0]\n", "\n", " # -------------------------------------------------------------\n", " # Encrypt result\n", " # -------------------------------------------------------------\n", "\n", " encrypted_result = encrypt(\n", " pil_to_bytes(result)\n", " )\n", "\n", " out_path = os.path.join(\n", " output_folders[worker_id],\n", " f\"edited_{global_idx:06d}.enc\"\n", " )\n", "\n", " with open(out_path, \"wb\") as f:\n", " f.write(encrypted_result)\n", "\n", " print(\n", " f\"โœ… Worker {worker_id} \"\n", " f\"(GPU {gpu_id}) \"\n", " f\"โ†’ {global_idx+1}/{len(input_files)}\"\n", " )\n", "\n", " del result\n", " del input_image\n", " del bg_image\n", "\n", " except Exception:\n", "\n", " print(\n", " f\"โŒ Worker {worker_id} \"\n", " f\"failed at index {global_idx}\"\n", " )\n", "\n", " print(traceback.format_exc())\n", "\n", " # -------------------------------------------------------------------------\n", " # Cleanup ZIP handles\n", " # -------------------------------------------------------------------------\n", "\n", " local_foreground_zip.close()\n", " local_background_zip.close()\n", "\n", " return True\n", "\n", "# =============================================================================\n", "# RUN WORKERS\n", "# =============================================================================\n", "\n", "chunk_size = (\n", " len(input_files) + len(pipes) - 1\n", ") // len(pipes)\n", "\n", "print(f\"\\n๐Ÿš€ Starting {len(pipes)} parallel workers...\")\n", "\n", "with ThreadPoolExecutor(max_workers=len(pipes)) as executor:\n", "\n", " futures = []\n", "\n", " for worker_id in range(len(pipes)):\n", "\n", " start = worker_id * chunk_size\n", "\n", " num = min(\n", " chunk_size,\n", " len(input_files) - start\n", " )\n", "\n", " if num > 0:\n", "\n", " futures.append(\n", "\n", " executor.submit(\n", " worker_thread,\n", " pipes[worker_id],\n", " gpu_ids[worker_id],\n", " start,\n", " num,\n", " worker_id\n", " )\n", "\n", " )\n", "\n", " for future in as_completed(futures):\n", " future.result()\n", "\n", "# =============================================================================\n", "# CLOSE MAIN ZIPS\n", "# =============================================================================\n", "\n", "foreground_zip.close()\n", "background_zip.close()\n", "\n", "print(\"\\n๐ŸŽ‰ All workers complete!\")" ], "metadata": { "id": "61W2bIhVfPy0" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "c23cc485" }, "source": [ "#@markdown CELL 5: Combine Encrypted Outputs into a Single ZIP\n", "import os\n", "import zipfile\n", "from pathlib import Path\n", "\n", "# Define the output zip file path\n", "combined_zip_path = \"/content/final_encrypted/results.zip\"\n", "\n", "# List the directories containing the encrypted output files\n", "output_dirs = [\n", " \"/content/edited_gpu0_pipe0\",\n", " \"/content/edited_gpu0_pipe1\"\n", "]\n", "\n", "all_encrypted_files = []\n", "for output_dir in output_dirs:\n", " if os.path.exists(output_dir):\n", " for f in os.listdir(output_dir):\n", " if f.endswith(\".enc\"):\n", " all_encrypted_files.append(os.path.join(output_dir, f))\n", "\n", "# Sort the files to ensure consistent enumeration (optional, but good practice)\n", "all_encrypted_files.sort()\n", "\n", "print(f\"Found {len(all_encrypted_files)} encrypted files to combine.\")\n", "\n", "if not all_encrypted_files:\n", " print(\"โš ๏ธ No encrypted files found to combine. Skipping zip creation.\")\n", "else:\n", " # Create a new zip file and add the encrypted files with new names\n", " with zipfile.ZipFile(combined_zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:\n", " for i, file_path in enumerate(all_encrypted_files):\n", " # New filename format: edited_000001.enc, edited_000002.enc, etc.\n", " new_file_name = f\"edited_{i+1:06d}.enc\"\n", " zf.write(file_path, arcname=new_file_name)\n", " print(f\" Added {os.path.basename(file_path)} as {new_file_name} to zip.\")\n", "\n", " print(f\"โœ… Successfully created combined zip file: {combined_zip_path}\")\n", " print(f\" Total files in zip: {len(all_encrypted_files)}\")" ], "execution_count": null, "outputs": [] } ] }