codeShare/lora-training-data
2785
1{"cells":[{"cell_type":"markdown","metadata":{"id":"HX3I-YqowwO9"},"source":["# Klein Edit"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"OXGhsHQwWue_"},"outputs":[],"source":["# =============================================================================\n","#@markdown # **CELL 1**: Mount Drive + HF auth\n","# =============================================================================\n","\n","from google.colab import drive, userdata\n","from huggingface_hub import login\n","import torch\n","import os\n","import gc\n","import shutil\n","\n","drive.mount('/content/drive')\n","\n","hf_token = userdata.get('HF_TOKEN')\n","if hf_token:\n"," login(token=hf_token)\n","else:\n"," print(\"⚠️ No HF_TOKEN found in secrets.\")\n","\n","print(\"🧹 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","\n","print(\"🔄 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","\n","print(\"✅ Cell 1 complete! RESTART RUNTIME now, then run Cell 2 → Cell 3 → Cell 4\")"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"a1wZv8PrXAe2"},"outputs":[],"source":["# =============================================================================\n","#@markdown # **CELL 2**: Fixed settings (resolution + options)\n","# =============================================================================\n","\n","resolution = '1024 x 1024 (Square)' #@param [\"1024 x 1024 (Square)\", \"512 x 1024 (Portrait)\", \"768 x 1024 (Slight Portrait)\", \"1536 x 1024 (Landscape)\", \"2048 x 1024 (Wide Landscape)\"] {type:\"string\"}\n","use_txt_prompts = True #@param {type:\"boolean\"}\n","debug = False #@param {type:\"boolean\"}\n","\n","res_map = {\n"," \"1024 x 1024 (Square)\": (1024, 1024),\n"," \"512 x 1024 (Portrait)\": (512, 1024),\n"," \"768 x 1024 (Slight Portrait)\": (768, 1024),\n"," \"1536 x 1024 (Landscape)\": (1536, 1024),\n"," \"2048 x 1024 (Wide Landscape)\": (2048, 1024)\n","}\n","target_width, target_height = res_map[resolution]\n","\n","print(\"✅ Cell 2 settings loaded\")\n","print(f\" Resolution: {target_width}×{target_height}\")\n","print(f\" Use .txt prompts: {use_txt_prompts}\")\n","print(f\" Debug mode: {debug}\")\n","print(\"\\nNow run Cell 3 (model load), then Cell 4 (inference)\")"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"WdNfo33NXNEa"},"outputs":[],"source":["# =============================================================================\n","#@markdown # **CELL 3**: Load QUANTIZED Flux2KleinPipeline + SDNQ optimizations\n","#(model stays in memory)\n","# =============================================================================\n","\n","import torch\n","import gc\n","from diffusers import Flux2KleinPipeline\n","import zipfile\n","import glob\n","import os\n","\n","print(\"📦 Installing sdnq...\")\n","!pip install -q sdnq\n","\n","from sdnq.common import use_torch_compile as triton_is_available\n","from sdnq.loader import apply_sdnq_options_to_model\n","\n","gc.collect()\n","torch.cuda.empty_cache()\n","\n","print(\"🔄 Loading QUANTIZED Flux2KleinPipeline (4-bit SDNQ + float16 for T4)...\")\n","pipe = Flux2KleinPipeline.from_pretrained(\n"," \"Disty0/FLUX.2-klein-4B-SDNQ-4bit-dynamic\",\n"," torch_dtype=torch.float16\n",")\n","\n","print(f\"✅ Pipeline loaded – VRAM: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")\n","\n","# ====================== Optimizations for T4 ======================\n","print(\"🔥 Applying SDNQ + T4 optimizations...\")\n","os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n","\n","if triton_is_available and torch.cuda.is_available():\n"," pipe.transformer = apply_sdnq_options_to_model(pipe.transformer, use_quantized_matmul=True)\n"," pipe.text_encoder = apply_sdnq_options_to_model(pipe.text_encoder, use_quantized_matmul=True)\n"," print(\" ✅ INT8 MatMul enabled\")\n","\n","pipe.enable_model_cpu_offload()\n","pipe.vae.enable_slicing()\n","pipe.vae.enable_tiling()\n","\n","gc.collect()\n","torch.cuda.empty_cache()\n","torch.cuda.reset_peak_memory_stats()\n","\n","print(\"✅ Cell 3 complete – model is ready!\")\n","print(\" Now go to Cell 4 and run it (you can change zip/prompt in Cell 4 and re-run Cell 4 only)\")"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","colab":{"background_save":true},"id":"q9I9wxghuUGP"},"outputs":[],"source":["# =============================================================================\n","#@markdown # **CELL 4**: MULTI-IMAGE INFERENCE (Klein + custom/gray reference)\n","# • Uses the **same zip** as Cell 4 but now does true multi-image editing\n","# • First image = your original (the one being edited)\n","# • Second image = your custom uploaded image OR a solid #181818 gray square (perfect for background replacement)\n","# • Re-run this cell only with new values\n","# • Auto-clears old data → processes → saves **only ZIP** to Drive\n","# =============================================================================\n","\n","edit_prompt = 'remove the background. the background is gray. ' #@param {type:\"string\"}\n","zip_path = '/content/drive/MyDrive/images.zip' #@param {type:\"string\"}\n","\n","import shutil\n","import datetime\n","import os\n","import glob\n","import zipfile\n","import gc\n","import torch\n","from PIL import Image\n","from google.colab import files\n","\n","output_folder = '/content/edited_images_multi' # ← separate temp folder so it doesn't clash with Cell 4\n","\n","print(\"🧹 Clearing old temporary folders so you can reuse with new zip/prompt...\")\n","if os.path.exists('/content/input_images'):\n"," shutil.rmtree('/content/input_images')\n","if os.path.exists(output_folder):\n"," shutil.rmtree(output_folder)\n","os.makedirs(output_folder, exist_ok=True)\n","\n","# ====================== Unzip + build prompt map ======================\n","print(f\"📦 Unzipping: {zip_path}...\")\n","with zipfile.ZipFile(zip_path, 'r') as z:\n"," z.extractall('/content/input_images')\n","\n","image_files = sorted(glob.glob('/content/input_images/*.*'))\n","image_files = [f for f in image_files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]\n","print(f\"Found {len(image_files)} images.\")\n","\n","prompt_map = {}\n","txt_count = 0\n","for img_path in image_files:\n"," base = os.path.splitext(os.path.basename(img_path))[0]\n"," txt_path = os.path.join('/content/input_images', f\"{base}.txt\")\n"," if use_txt_prompts and os.path.exists(txt_path):\n"," with open(txt_path, 'r', encoding='utf-8') as f:\n"," prompt_map[img_path] = f.read().strip()\n"," txt_count += 1\n"," else:\n"," prompt_map[img_path] = None\n","\n","print(f\"Found {txt_count} matching .txt files\")\n","\n","# ====================== Create or load the second reference image (once) ======================\n","print(\"Please upload your reference image now, or skip to use a gray background.\")\n","uploaded_files = files.upload()\n","\n","uploaded_reference_filepath = None\n","if uploaded_files:\n"," uploaded_reference_filepath = list(uploaded_files.keys())[0]\n"," print(f\"✅ Uploaded reference image: {uploaded_reference_filepath}\")\n"," # Save the uploaded file to a temporary location to be opened by PIL\n"," with open(uploaded_reference_filepath, 'wb') as f:\n"," f.write(uploaded_files[uploaded_reference_filepath])\n","\n","reference_image_to_pair = None\n","if uploaded_reference_filepath: # Check if a file was actually uploaded\n"," print(\"📁 Using uploaded reference image...\")\n"," try:\n"," # Open the temporary file\n"," reference_image_to_pair = Image.open(uploaded_reference_filepath).convert(\"RGB\")\n"," #if reference_image_to_pair.size != (target_width, target_height):\n"," # print(f\" ⚠️ Resizing uploaded reference image from {reference_image_to_pair.size} to {target_width}x{target_height}\")\n"," # reference_image_to_pair = reference_image_to_pair.resize((target_width, target_height), Image.LANCZOS)\n"," except Exception as e:\n"," print(f\" ❌ Error loading uploaded image: {e}. Falling back to gray reference.\")\n"," reference_image_to_pair = Image.new(\"RGB\", (target_width, target_height), \"#181818\")\n","else:\n"," print(f\"🟩 No reference image uploaded. Creating {target_width}x{target_height} #181818 solid gray reference image...\")\n"," reference_image_to_pair = Image.new(\"RGB\", (target_width, target_height), \"#181818\")\n","\n","# Clean up the temporarily saved file if it was uploaded\n","if uploaded_reference_filepath and os.path.exists(uploaded_reference_filepath):\n"," os.remove(uploaded_reference_filepath)\n","\n","\n","# ====================== FULL BATCH MULTI-IMAGE INFERENCE ======================\n","print(f\"\\n🚀 Starting batch MULTI-IMAGE edit on {len(image_files)} images...\")\n","\n","for i, img_path in enumerate(image_files):\n"," filename = os.path.basename(img_path)\n"," print(f\"[{i+1}/{len(image_files)}] {filename} → + custom reference\") # Updated message\n","\n"," gc.collect()\n"," torch.cuda.empty_cache()\n","\n"," input_image = Image.open(img_path).convert(\"RGB\")\n","\n"," # === MULTI-IMAGE LIST: base image first, the chosen reference second ===\n"," reference_images = [input_image, reference_image_to_pair] # Use the chosen reference\n","\n"," current_prompt = prompt_map[img_path] if use_txt_prompts and prompt_map[img_path] is not None else edit_prompt\n","\n"," result = pipe(\n"," prompt=current_prompt,\n"," image=reference_images, # ← KEY: list of images for Klein multi-reference\n"," height=target_height,\n"," width=target_width,\n"," guidance_scale=1.0,\n"," num_inference_steps=4,\n"," generator=torch.Generator(\"cuda\").manual_seed(42),\n"," output_type=\"pil\",\n"," ).images[0]\n","\n"," output_filename = f\"multi_edited_{filename}\"\n"," output_path = os.path.join(output_folder, output_filename)\n"," result.save(output_path)\n","\n"," if debug:\n"," print(f\" → Saved {output_filename} (multi-ref with custom/gray image)\") # Updated message\n","\n","print(\"\\n🎉 MULTI-IMAGE BATCH COMPLETE!\")\n","\n","# ====================== AUTO ZIP TO DRIVE (unique name) ======================\n","timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n","drive_zip_path = f\"/content/drive/MyDrive/multi_edited_images_{timestamp}.zip\"\n","\n","print(f\"📦 Creating zip → {drive_zip_path}\")\n","shutil.make_archive(drive_zip_path.replace('.zip', ''), 'zip', output_folder)\n","\n","print(\"✅ ALL DONE!\")\n","print(f\"Zip file saved to your Google Drive:\")\n","print(drive_zip_path)"]},{"cell_type":"markdown","metadata":{"id":"11YbjJrYStFb"},"source":["# Quantized lora training (unfinished)"]},{"cell_type":"code","source":["#@markdown # **CELL 1**: Mount Drive + HF auth\n","from google.colab import drive, userdata\n","from huggingface_hub import login\n","import torch\n","import os\n","import gc\n","\n","drive.mount('/content/drive')\n","\n","hf_token = userdata.get('HF_TOKEN')\n","if hf_token:\n"," login(token=hf_token)\n"," print(\"✅ HF login successful – gated model download enabled\")\n","else:\n"," print(\"⚠️ No HF_TOKEN found in secrets. FLUX.2-klein-4B download may fail.\")\n","\n","print(\"🧹 Clearing GPU memory...\")\n","torch.cuda.empty_cache()\n","gc.collect()\n","\n","print(\"✅ Cell 1 complete! Now run Cell 2A\")"],"metadata":{"cellView":"form","id":"9mZkQOvtx-7k","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1776737911591,"user_tz":-120,"elapsed":40607,"user":{"displayName":"fukU Google","userId":"02763165356193834046"}},"outputId":"a0cdaee3-e825-4a4e-bd4c-275c6c2c744d"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Mounted at /content/drive\n","✅ HF login successful – gated model download enabled\n","🧹 Clearing GPU memory...\n","✅ Cell 1 complete! Now run Cell 2A\n"]}]},{"cell_type":"code","source":["#@markdown # **CELL 2A**: Environment Setup – Clean + FLUX requirements (FIXED pipeline import)\n","\n","import os\n","\n","print(\"🧹 Resetting working directory...\")\n","%cd /content\n","\n","# Remove old diffusers completely\n","print(\"🧹 Removing any existing diffusers...\")\n","!pip uninstall -y diffusers peft accelerate transformers bitsandbytes -q\n","!rm -rf /content/diffusers /usr/local/lib/python*/dist-packages/diffusers*\n","\n","print(\"🔄 Installing latest diffusers from GitHub main (this fixes Flux2KleinPipeline)...\")\n","!pip install -q git+https://github.com/huggingface/diffusers.git --force-reinstall --no-deps\n","\n","print(\"🔄 Cloning training scripts (we only need the dreambooth folder)...\")\n","!git clone https://github.com/huggingface/diffusers.git --depth=1 /content/diffusers\n","\n","# FLUX-specific requirements\n","print(\"🔄 Installing FLUX requirements...\")\n","!pip install -q -r /content/diffusers/examples/dreambooth/requirements_flux.txt --no-deps\n","\n","# bitsandbytes for QLoRA\n","print(\"🔄 Installing bitsandbytes...\")\n","!pip install -q bitsandbytes\n","\n","print(\"🔄 Installing dreambooth script requirements...\")\n","%cd /content/diffusers/examples/dreambooth\n","!pip install -q -r requirements.txt --no-deps\n","!pip install -q --upgrade peft\n","\n","print(\"\\n\" + \"=\"*70)\n","print(\"✅ DEBUG: Package versions\")\n","!pip list | grep -E 'diffusers|peft|accelerate|transformers|torch|bitsandbytes'\n","print(\"=\"*70)\n","\n","import diffusers\n","print(f\"✅ diffusers version: {diffusers.__version__}\")\n","print(\"✅ Cell 2A complete!\")\n","\n","print(\" → **Restart runtime now** (Runtime → Restart session)\")\n","print(\" → Then run the new Cell 2B\")"],"metadata":{"cellView":"form","id":"wjhleprr0Xj6","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1776737989210,"user_tz":-120,"elapsed":77622,"user":{"displayName":"fukU Google","userId":"02763165356193834046"}},"outputId":"09d4841e-a5fc-4251-b2ee-2dbbe4313f31"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["🧹 Resetting working directory...\n","/content\n","🧹 Removing any existing diffusers...\n","\u001b[33mWARNING: Skipping bitsandbytes as it is not installed.\u001b[0m\u001b[33m\n","\u001b[0m🔄 Installing latest diffusers from GitHub main (this fixes Flux2KleinPipeline)...\n"," Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n"," Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n"," Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"," Building wheel for diffusers (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n","🔄 Cloning training scripts (we only need the dreambooth folder)...\n","Cloning into '/content/diffusers'...\n","remote: Enumerating objects: 2809, done.\u001b[K\n","remote: Counting objects: 100% (2809/2809), done.\u001b[K\n","remote: Compressing objects: 100% (1947/1947), done.\u001b[K\n","remote: Total 2809 (delta 1116), reused 1525 (delta 833), pack-reused 0 (from 0)\u001b[K\n","Receiving objects: 100% (2809/2809), 9.01 MiB | 16.56 MiB/s, done.\n","Resolving deltas: 100% (1116/1116), done.\n","🔄 Installing FLUX requirements...\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m383.7/383.7 kB\u001b[0m \u001b[31m20.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.2/10.2 MB\u001b[0m \u001b[31m108.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.8/44.8 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m680.7/680.7 kB\u001b[0m \u001b[31m51.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25h🔄 Installing bitsandbytes...\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.7/60.7 MB\u001b[0m \u001b[31m12.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25h🔄 Installing dreambooth script requirements...\n","/content/diffusers/examples/dreambooth\n","\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m168.3/168.3 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25h\n","======================================================================\n","✅ DEBUG: Package versions\n","accelerate 1.13.0\n","bitsandbytes 0.49.2\n","diffusers 0.38.0.dev0\n","peft 0.19.1\n","sentence-transformers 5.4.0\n","torch 2.10.0+cu128\n","torchao 0.10.0\n","torchaudio 2.10.0+cu128\n","torchcodec 0.10.0+cu128\n","torchdata 0.11.0\n","torchsummary 1.5.1\n","torchtune 0.6.1\n","torchvision 0.25.0+cu128\n","transformers 5.5.4\n","======================================================================\n","✅ diffusers version: 0.38.0.dev0\n","✅ Cell 2A complete!\n"," → **Restart runtime now** (Runtime → Restart session)\n"," → Then run the new Cell 2B\n"]}]},{"cell_type":"code","source":["#@markdown # **CELL 2B**: Verify diffusers + Flux2KleinPipeline (now fixed)\n","\n","import torch\n","import gc\n","import sys\n","import importlib\n","\n","print(\"🔄 Verifying environment after restart...\")\n","\n","# No more path hacks needed — git install puts everything in the right place\n","importlib.invalidate_caches()\n","\n","import diffusers\n","print(f\"✅ diffusers imported from: {diffusers.__file__}\")\n","print(f\" version: {diffusers.__version__}\")\n","\n","# Test the pipeline that was failing\n","print(\"\\n🔍 Testing Flux2KleinPipeline import...\")\n","from diffusers import Flux2KleinPipeline\n","print(\"✅ Flux2KleinPipeline imported successfully! 🎉\")\n","\n","print(\"\\n\" + \"=\"*80)\n","print(\"✅ ENVIRONMENT IS NOW FULLY READY FOR QLoRA TRAINING\")\n","print(\" → You can safely run Cell 3 → Cell 4\")\n","print(\"=\"*80)\n","\n","torch.cuda.empty_cache()\n","gc.collect()"],"metadata":{"cellView":"form","id":"9wZsOy760Yha","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1776738025559,"user_tz":-120,"elapsed":25290,"user":{"displayName":"fukU Google","userId":"02763165356193834046"}},"outputId":"e2845a62-f332-4474-a0c2-b752380006ad"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["🔄 Verifying environment after restart...\n","✅ diffusers imported from: /usr/local/lib/python3.12/dist-packages/diffusers/__init__.py\n"," version: 0.38.0.dev0\n","\n","🔍 Testing Flux2KleinPipeline import...\n"]},{"output_type":"stream","name":"stderr","text":["Flax classes are deprecated and will be removed in Diffusers v1.0.0. We recommend migrating to PyTorch classes or pinning your version of Diffusers.\n","Flax classes are deprecated and will be removed in Diffusers v1.0.0. We recommend migrating to PyTorch classes or pinning your version of Diffusers.\n"]},{"output_type":"stream","name":"stdout","text":["✅ Flux2KleinPipeline imported successfully! 🎉\n","\n","================================================================================\n","✅ ENVIRONMENT IS NOW FULLY READY FOR QLoRA TRAINING\n"," → You can safely run Cell 3 → Cell 4\n","================================================================================\n"]},{"output_type":"execute_result","data":{"text/plain":["0"]},"metadata":{},"execution_count":3}]},{"cell_type":"code","source":["#@markdown # **CELL 3**: Prepare environment + dataset\n","import os\n","import shutil\n","import glob\n","import zipfile\n","from PIL import Image\n","import torch\n","import gc\n","\n","# ====================== SETTINGS ======================\n","training_zip_path = '/content/drive/MyDrive/morbids/morbid4.zip' #@param {type:\"string\"}\n","max_images = 2 #@param {type:\"integer\"} # Increase on Pro+\n","\n","train_data_dir = \"/content/colorization_train_pairs\"\n","output_dir = \"/content/flux_klein_colorize_lora\"\n","\n","print(\"🧹 Clearing old folders...\")\n","if os.path.exists(train_data_dir):\n"," shutil.rmtree(train_data_dir)\n","if os.path.exists(output_dir):\n"," shutil.rmtree(output_dir)\n","os.makedirs(train_data_dir, exist_ok=True)\n","\n","# ====================== LOAD DATASET ======================\n","print(f\"📦 Unzipping: {training_zip_path}\")\n","with zipfile.ZipFile(training_zip_path, 'r') as z:\n"," z.extractall('/content/raw_training_images')\n","\n","image_files = sorted(glob.glob('/content/raw_training_images/*.*'))\n","image_files = [f for f in image_files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))][:max_images]\n","\n","print(f\"Found {len(image_files)} images → preparing for colorization training...\")\n","\n","for i, img_path in enumerate(image_files):\n"," base = os.path.splitext(os.path.basename(img_path))[0]\n"," color_img = Image.open(img_path).convert(\"RGB\")\n"," color_img.save(os.path.join(train_data_dir, f\"{base}.png\"))\n","\n"," # Same prompt for every image\n"," with open(os.path.join(train_data_dir, f\"{base}.txt\"), \"w\", encoding=\"utf-8\") as f:\n"," f.write(\"colorize this image\")\n","\n","print(f\"✅ Dataset ready – {len(image_files)} examples in {train_data_dir}\")\n","\n","# ====================== VRAM OPTIMIZATIONS ======================\n","os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n","gc.collect()\n","torch.cuda.empty_cache()\n","\n","print(\"✅ Cell 3 complete! Ready for training.\")"],"metadata":{"id":"6nUSv-atykR-","colab":{"base_uri":"https://localhost:8080/"},"cellView":"form","executionInfo":{"status":"ok","timestamp":1776738060571,"user_tz":-120,"elapsed":7239,"user":{"displayName":"fukU Google","userId":"02763165356193834046"}},"outputId":"d2aa52b8-3019-4710-a15f-b92f060a29d2"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["🧹 Clearing old folders...\n","📦 Unzipping: /content/drive/MyDrive/morbids/morbid4.zip\n","Found 2 images → preparing for colorization training...\n","✅ Dataset ready – 2 examples in /content/colorization_train_pairs\n","✅ Cell 3 complete! Ready for training.\n"]}]},{"cell_type":"code","source":["#@markdown # **CELL 4A**: Setup (bnb config + clean dataset folder)\n","\n","import torch\n","import gc\n","import json\n","import os\n","import glob\n","\n","print(\"🚀 Preparing quantized QLoRA training...\")\n","\n","# 4-bit NF4 config\n","bnb_config_path = \"/content/bnb_config.json\"\n","bnb_config = {\n"," \"load_in_4bit\": True,\n"," \"bnb_4bit_quant_type\": \"nf4\",\n"," \"bnb_4bit_compute_dtype\": \"float16\",\n"," \"bnb_4bit_use_double_quant\": True\n","}\n","with open(bnb_config_path, \"w\") as f:\n"," json.dump(bnb_config, f)\n","print(f\"✅ 4-bit config saved → {bnb_config_path}\")\n","\n","# ====================== CRITICAL FIX: Remove .txt files ======================\n","print(\"🧹 Removing all .txt files from training folder (script only wants images)...\")\n","train_data_dir = \"/content/colorization_train_pairs\"\n","txt_files = glob.glob(os.path.join(train_data_dir, \"*.txt\"))\n","for txt in txt_files:\n"," os.remove(txt)\n"," print(f\" Deleted: {os.path.basename(txt)}\")\n","\n","print(f\"✅ Kept only {len(glob.glob(os.path.join(train_data_dir, '*.png')))} image files\")\n","print(\" → Training folder is now clean for --instance_prompt mode\")\n","\n","# Optional: clear memory before launch\n","gc.collect()\n","torch.cuda.empty_cache()\n","\n","print(\"\\n✅ Cell 4A complete! Now run Cell 4B to start training\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"RLpLJZHt6QXI","executionInfo":{"status":"ok","timestamp":1776739239157,"user_tz":-120,"elapsed":696,"user":{"displayName":"fukU Google","userId":"02763165356193834046"}},"outputId":"4a1a96d3-4ef0-4f55-cdaf-775e608310bd"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["🚀 Preparing quantized QLoRA training...\n","✅ 4-bit config saved → /content/bnb_config.json\n","🧹 Removing all .txt files from training folder (script only wants images)...\n"," Deleted: 53537172992_145a9daed1_k.txt\n"," Deleted: 53537173032_18aa14e83b_k.txt\n","✅ Kept only 2 image files\n"," → Training folder is now clean for --instance_prompt mode\n","\n","✅ Cell 4A complete! Now run Cell 4B to start training\n"]}]},{"cell_type":"code","source":["#@markdown # **CELL 4B**: Quantized QLoRA Training – **Sliders + Tiny Dataset + LOSS LOGGING**\n","\n","import datetime\n","import os\n","import re\n","\n","print(\"🚀 Starting **quantized QLoRA** on FLUX.2-klein-4B (T4-safe + sliders + loss logging)...\")\n","\n","# ========================== TRAINING PARAMETERS ==========================\n","max_train_steps = 100 #@param {type:\"slider\", min:100, max:2000, step:50}\n","learning_rate = 1e-4 #@param {type:\"slider\", min:5e-5, max:5e-4, step:5e-5, format:\"0.0e\"}\n","rank = 16 #@param {type:\"slider\", min:8, max:64, step:8}\n","gradient_accumulation_steps = 4 #@param {type:\"slider\", min:4, max:16, step:4}\n","lr_warmup_steps = 10 #@param {type:\"slider\", min:0, max:200, step:10}\n","lora_alpha = 32 #@param {type:\"slider\", min:16, max:64, step:8}\n","\n","print(f\" • max_train_steps = {max_train_steps}\")\n","print(f\" • learning_rate = {learning_rate}\")\n","print(f\" • rank = {rank}\")\n","print(f\" • gradient_accumulation = {gradient_accumulation_steps}\")\n","print(f\" • lr_warmup_steps = {lr_warmup_steps}\")\n","print(f\" • lora_alpha = {lora_alpha}\")\n","print(\" → Very low settings for tiny 2-image dataset\")\n","\n","# ====================== TRAINING (with live output + full log) ======================\n","print(\"\\n🔥 Launching training... (live progress will be shown below)\")\n","\n","!accelerate launch train_dreambooth_lora_flux2_klein.py \\\n"," --pretrained_model_name_or_path=\"black-forest-labs/FLUX.2-klein-4B\" \\\n"," --instance_data_dir=\"/content/colorization_train_pairs\" \\\n"," --instance_prompt=\"colorize this image\" \\\n"," --resolution=1024 \\\n"," --train_batch_size=1 \\\n"," --gradient_accumulation_steps={gradient_accumulation_steps} \\\n"," --gradient_checkpointing \\\n"," --cache_latents \\\n"," --learning_rate={learning_rate} \\\n"," --lr_scheduler=\"constant\" \\\n"," --lr_warmup_steps={lr_warmup_steps} \\\n"," --max_train_steps={max_train_steps} \\\n"," --use_8bit_adam \\\n"," --mixed_precision=\"fp16\" \\\n"," --rank={rank} \\\n"," --lora_alpha={lora_alpha} \\\n"," --output_dir=\"/content/flux_klein_colorize_lora\" \\\n"," --report_to=\"tensorboard\" \\\n"," --seed=42 \\\n"," --bnb_quantization_config_path=\"/content/bnb_config.json\" \\\n"," --offload \\\n"," --skip_final_inference 2>&1 | tee training_log.txt\n","\n","# ====================== PARSE AND PRINT CLEAN LOSS SUMMARY ======================\n","print(\"\\n\" + \"=\"*70)\n","print(\"📊 TRAINING LOSS SUMMARY (copy-paste this part to me)\")\n","print(\"=\"*70)\n","\n","if os.path.exists(\"training_log.txt\"):\n"," with open(\"training_log.txt\", \"r\", encoding=\"utf-8\") as f:\n"," log_content = f.read()\n","\n"," # Extract all loss= values from tqdm progress lines\n"," loss_matches = re.findall(r'loss=([\\d.]+)', log_content)\n","\n"," if loss_matches:\n"," losses = [float(x) for x in loss_matches]\n"," total_steps = len(losses)\n","\n"," print(f\"✅ Training completed {total_steps} steps\")\n"," print(f\" Average loss: {sum(losses)/total_steps:.4f}\")\n"," print(f\" Final loss: {losses[-1]:.4f}\")\n","\n"," print(\"\\nLast 15 losses:\")\n"," for i, loss in enumerate(losses[-15:]):\n"," step_num = total_steps - 15 + i + 1\n"," print(f\" Step {step_num:4d} → loss = {loss:.4f}\")\n","\n"," # Optional: save full log to Drive\n"," timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n"," log_zip = f\"/content/drive/MyDrive/flux_klein_colorize_lora_log_{timestamp}.txt\"\n"," with open(log_zip, \"w\", encoding=\"utf-8\") as f:\n"," f.write(log_content)\n"," print(f\"\\n📁 Full training log saved to Drive: {log_zip}\")\n"," else:\n"," print(\"⚠️ No loss values found in log (training may have crashed early)\")\n","else:\n"," print(\"⚠️ training_log.txt not found\")\n","\n","print(\"\\n🎉 TRAINING COMPLETE!\")\n","print(f\"LoRA saved to: /content/flux_klein_colorize_lora\")\n","\n","# Save LoRA to Drive\n","timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n","drive_zip = f\"/content/drive/MyDrive/flux_klein_colorize_lora_{timestamp}.zip\"\n","!zip -r -q \"{drive_zip}\" /content/flux_klein_colorize_lora\n","print(f\"✅ LoRA ZIP saved to Drive: {drive_zip}\")"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"4r51ADIT9M6s","outputId":"dc26d6e8-2cc9-4645-ff01-febc95a62c39"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["🚀 Starting **quantized QLoRA** on FLUX.2-klein-4B (T4-safe + sliders + loss logging)...\n"," • max_train_steps = 100\n"," • learning_rate = 0.0001\n"," • rank = 16\n"," • gradient_accumulation = 4\n"," • lr_warmup_steps = 10\n"," • lora_alpha = 32\n"," → Very low settings for tiny 2-image dataset\n","\n","🔥 Launching training... (live progress will be shown below)\n","The following values were not passed to `accelerate launch` and had defaults used instead:\n","\t`--num_processes` was set to a value of `1`\n","\t`--num_machines` was set to a value of `1`\n","\t`--mixed_precision` was set to a value of `'no'`\n","\t`--dynamo_backend` was set to a value of `'no'`\n","To avoid this warning pass in values for each of the problematic parameters or run `accelerate config`.\n","Flax classes are deprecated and will be removed in Diffusers v1.0.0. We recommend migrating to PyTorch classes or pinning your version of Diffusers.\n","Flax classes are deprecated and will be removed in Diffusers v1.0.0. We recommend migrating to PyTorch classes or pinning your version of Diffusers.\n","/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_validators.py:205: UserWarning: The `local_dir_use_symlinks` argument is deprecated and ignored in `hf_hub_download`. Downloading to a local directory does not use symlinks anymore.\n"," warnings.warn(\n","{'decoder_block_out_channels'} was not found in config. Values will be initialized to default values.\n","All model checkpoint weights were used when initializing AutoencoderKLFlux2.\n","\n","All the weights of AutoencoderKLFlux2 were initialized from the model checkpoint at black-forest-labs/FLUX.2-klein-4B.\n","If your task is similar to the task the model of the checkpoint was trained on, you can already use AutoencoderKLFlux2 for predictions without further training.\n","The device_map was not initialized. Setting device_map to {: {current_device}}. If you want to use the model for inference, please set device_map ='auto' \n","Instantiating Flux2Transformer2DModel model under default dtype torch.float16.\n","All model checkpoint weights were used when initializing Flux2Transformer2DModel.\n","\n","All the weights of Flux2Transformer2DModel were initialized from the model checkpoint at black-forest-labs/FLUX.2-klein-4B.\n","If your task is similar to the task the model of the checkpoint was trained on, you can already use Flux2Transformer2DModel for predictions without further training.\n","Fetching 2 files: 100%|██████████| 2/2 [00:00<00:00, 1596.01it/s]\n","Loading weights: 100%|██████████| 398/398 [00:00<00:00, 419.91it/s]\n","Loading pipeline components...: 100%|██████████| 2/2 [00:00<00:00, 4984.32it/s]\n","Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n","Caching latents: 100%|██████████| 2/2 [00:02<00:00, 1.08s/it]\n","Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n","2026-04-21 02:55:58.822715: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n","WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n","E0000 00:00:1776740159.066393 17109 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n","E0000 00:00:1776740159.129461 17109 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n","W0000 00:00:1776740159.645169 17109 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n","W0000 00:00:1776740159.645215 17109 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n","W0000 00:00:1776740159.645220 17109 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n","W0000 00:00:1776740159.645222 17109 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n","Steps: 94%|█████████▍| 94/100 [32:48<02:04, 20.82s/it, loss=0.734, lr=0.0001]"]}]}],"metadata":{"accelerator":"GPU","colab":{"collapsed_sections":["HX3I-YqowwO9"],"gpuType":"T4","provenance":[{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/klein_edit.ipynb","timestamp":1776702729526},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/vertical_slice_prepper.ipynb","timestamp":1776687886662},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/vertical_slice_prepper.ipynb","timestamp":1776366149549},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/lora_vertical_slice_dataset_creator.ipynb","timestamp":1776287741995},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/lora_vertical_slice_dataset_creator.ipynb","timestamp":1776178739426},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/lora_vertical_slice_dataset_creator.ipynb","timestamp":1776027716448},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/lora_vertical_slice_dataset_creator.ipynb","timestamp":1773663661932},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/civit_caption_prepper.ipynb","timestamp":1773663290922},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/civit_caption_prepper.ipynb","timestamp":1773264797996},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/civit_caption_prepper.ipynb","timestamp":1773163850245},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/civit_caption_prepper.ipynb","timestamp":1773090196076},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/civit_caption_prepper.ipynb","timestamp":1773089575687},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/civit_caption_prepper.ipynb","timestamp":1773080355474},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/Drive to WebP.ipynb","timestamp":1772998638620},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/Drive to WebP.ipynb","timestamp":1763646205520},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/Drive to WebP.ipynb","timestamp":1760993725927},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/YT-playlist-to-mp3.ipynb","timestamp":1760450712160},{"file_id":"https://huggingface.co/datasets/codeShare/lora-training-data/blob/main/YT-playlist-to-mp3.ipynb","timestamp":1756712618300},{"file_id":"https://huggingface.co/codeShare/JupyterNotebooks/blob/main/YT-playlist-to-mp3.ipynb","timestamp":1747490904984},{"file_id":"https://huggingface.co/codeShare/JupyterNotebooks/blob/main/YT-playlist-to-mp3.ipynb","timestamp":1740037333374},{"file_id":"https://huggingface.co/codeShare/JupyterNotebooks/blob/main/YT-playlist-to-mp3.ipynb","timestamp":1736477078136},{"file_id":"https://huggingface.co/codeShare/JupyterNotebooks/blob/main/YT-playlist-to-mp3.ipynb","timestamp":1725365086834}]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}