{ "cells": [ { "cell_type": "markdown", "source": [ "# create SDNQ" ], "metadata": { "id": "9vXxafECsKC_" } }, { "cell_type": "code", "source": [ "#@markdown # CELL 1 — Setup environment + login\n", "\n", "from google.colab import drive, userdata\n", "from huggingface_hub import login\n", "import torch, os, 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", "else:\n", " raise ValueError(\"HF_TOKEN not found in Google Colab secrets\")\n", "\n", "print(\"✅ Logged into Hugging Face\")" ], "metadata": { "id": "XzgpnXJhsHg-" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@markdown # CELL 2 — Install dependencies\n", "\n", "!pip install -q safetensors huggingface_hub diffusers transformers accelerate\n", "\n", "print(\"✅ Dependencies installed\")" ], "metadata": { "id": "XXld_BQ1sU0v" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "import torch\n", "from transformers import BitsAndBytesConfig\n", "from diffusers import DiffusionPipeline\n", "\n", "# ============================================================\n", "# BNB CONFIG\n", "# ============================================================\n", "\n", "bnb_config = BitsAndBytesConfig(\n", " load_in_4bit=True,\n", " bnb_4bit_quant_type=\"nf4\",\n", " bnb_4bit_compute_dtype=torch.bfloat16,\n", " bnb_4bit_use_double_quant=True,\n", ")\n", "\n", "# ============================================================\n", "# LOAD PIPELINE\n", "# ============================================================\n", "\n", "pipe = DiffusionPipeline.from_pretrained(\n", " \"black-forest-labs/FLUX.2-klein-9B\",\n", " # Removed: quantization_config=bnb_config,\n", " torch_dtype=torch.bfloat16,\n", "\n", " # Pass individual quantization parameters directly\n", " load_in_4bit=bnb_config.load_in_4bit,\n", " bnb_4bit_quant_type=bnb_config.bnb_4bit_quant_type,\n", " bnb_4bit_compute_dtype=bnb_config.bnb_4bit_compute_dtype,\n", " bnb_4bit_use_double_quant=bnb_config.bnb_4bit_use_double_quant,\n", "\n", " # critical:\n", " device_map=\"cpu\", # Changed 'auto' to 'cpu'\n", "\n", " # enables CPU offload folder\n", " offload_folder=\"/content/offload\",\n", "\n", " # low RAM loading\n", " low_cpu_mem_usage=True,\n", ")\n", "\n", "# ============================================================\n", "# ENABLE FULL CPU OFFLOAD\n", "# ============================================================\n", "\n", "pipe.enable_model_cpu_offload()\n", "\n", "print(\"✅ Loaded with bitsandbytes + CPU offload\")" ], "metadata": { "id": "i9MbakNlsckw" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "import gc , torch\n", "gc.collect()\n", "torch.cuda.empty_cache()" ], "metadata": { "id": "rQ5cr5ZsUQqG" }, "execution_count": 4, "outputs": [] }, { "cell_type": "code", "source": [ "!pip install bitsandbytes\n", "import torch\n", "import gc\n", "import bitsandbytes as bnb\n", "\n", "from torch import nn\n", "\n", "# ============================================================\n", "# DEQUANTIZE BNB MODEL\n", "# ============================================================\n", "\n", "def dequantize_bnb_linear(layer):\n", " \"\"\"\n", " Convert Linear4bit -> torch.nn.Linear\n", " \"\"\"\n", "\n", " # Reconstruct full precision weights\n", " weight = layer.weight.dequantize()\n", "\n", " new_layer = nn.Linear(\n", " layer.in_features,\n", " layer.out_features,\n", " bias=layer.bias is not None,\n", " dtype=weight.dtype,\n", " device=\"cpu\",\n", " )\n", "\n", " new_layer.weight.data.copy_(weight.cpu())\n", "\n", " if layer.bias is not None:\n", " new_layer.bias.data.copy_(layer.bias.data.cpu())\n", "\n", " return new_layer\n", "\n", "\n", "def recursively_dequantize(module):\n", " \"\"\"\n", " Recursively replace all BNB 4bit layers.\n", " \"\"\"\n", "\n", " for name, child in list(module.named_children()):\n", "\n", " # Replace Linear4bit\n", " if isinstance(child, bnb.nn.Linear4bit):\n", " print(f\"Dequantizing: {name}\")\n", "\n", " setattr(\n", " module,\n", " name,\n", " dequantize_bnb_linear(child)\n", " )\n", "\n", " else:\n", " recursively_dequantize(child)\n", "\n", "\n", "# ============================================================\n", "# RUN DEQUANTIZATION\n", "# ============================================================\n", "\n", "pipe.to(\"cpu\")\n", "\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "# Apply dequantization to the individual sub-modules within the pipeline\n", "print(\"Dequantizing pipe.transformer...\")\n", "recursively_dequantize(pipe.transformer)\n", "print(\"Dequantizing pipe.text_encoder...\")\n", "recursively_dequantize(pipe.text_encoder)\n", "print(\"Dequantizing pipe.vae...\")\n", "recursively_dequantize(pipe.vae)\n", "\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "print(\"✅ Fully dequantized back to FP weights\")" ], "metadata": { "id": "GbZh_p6zSeTy" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#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", "!pip install -q --upgrade huggingface_hub transformers accelerate diffusers\n", "!pip install -q sdnq\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!\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "aGH03dqUQcsQ", "outputId": "52fdb852-4c99-4e2c-f13e-7bfe0e62f511" }, "execution_count": 6, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "✅ Cell 1 complete!\n" ] } ] }, { "cell_type": "code", "source": [ "from sdnq.common import accepted_weight_dtypes\n", "accepted_weight_dtypes" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "EpL4Kek7AbjZ", "outputId": "9c46f98d-cbd0-498f-bb94-33b6860f49c6" }, "execution_count": null, "outputs": [ { "output_type": "execute_result", "data": { "text/plain": [ "{'bf16',\n", " 'bfloat16',\n", " 'bool',\n", " 'float10_e1m8fn',\n", " 'float10_e1m9fnu',\n", " 'float10_e2m7fn',\n", " 'float10_e2m8fnu',\n", " 'float10_e3m6fn',\n", " 'float10_e3m7fnu',\n", " 'float10_e4m5fn',\n", " 'float10_e4m6fnu',\n", " 'float10_e5m4fn',\n", " 'float10_e5m5fnu',\n", " 'float11_e1m10fnu',\n", " 'float11_e1m9fn',\n", " 'float11_e2m8fn',\n", " 'float11_e2m9fnu',\n", " 'float11_e3m7fn',\n", " 'float11_e3m8fnu',\n", " 'float11_e4m6fn',\n", " 'float11_e4m7fnu',\n", " 'float11_e5m5fn',\n", " 'float11_e5m6fnu',\n", " 'float12_e1m10fn',\n", " 'float12_e1m11fnu',\n", " 'float12_e2m10fnu',\n", " 'float12_e2m9fn',\n", " 'float12_e3m8fn',\n", " 'float12_e3m9fnu',\n", " 'float12_e4m7fn',\n", " 'float12_e4m8fnu',\n", " 'float12_e5m6fn',\n", " 'float12_e5m7fnu',\n", " 'float13_e1m11fn',\n", " 'float13_e1m12fnu',\n", " 'float13_e2m10fn',\n", " 'float13_e2m11fnu',\n", " 'float13_e3m10fnu',\n", " 'float13_e3m9fn',\n", " 'float13_e4m8fn',\n", " 'float13_e4m9fnu',\n", " 'float13_e5m7fn',\n", " 'float13_e5m8fnu',\n", " 'float14_e1m12fn',\n", " 'float14_e1m13fnu',\n", " 'float14_e2m11fn',\n", " 'float14_e2m12fnu',\n", " 'float14_e3m10fn',\n", " 'float14_e3m11fnu',\n", " 'float14_e4m10fnu',\n", " 'float14_e4m9fn',\n", " 'float14_e5m8fn',\n", " 'float14_e5m9fnu',\n", " 'float15_e1m13fn',\n", " 'float15_e1m14fnu',\n", " 'float15_e2m12fn',\n", " 'float15_e2m13fnu',\n", " 'float15_e3m11fn',\n", " 'float15_e3m12fnu',\n", " 'float15_e4m10fn',\n", " 'float15_e4m11fnu',\n", " 'float15_e5m10fnu',\n", " 'float15_e5m9fn',\n", " 'float16',\n", " 'float16_e1m14fn',\n", " 'float16_e1m15fnu',\n", " 'float16_e2m13fn',\n", " 'float16_e2m14fnu',\n", " 'float16_e3m12fn',\n", " 'float16_e3m13fnu',\n", " 'float16_e4m11fn',\n", " 'float16_e4m12fnu',\n", " 'float16_e5m10fn',\n", " 'float16_e5m11fnu',\n", " 'float1_e1m0fnu',\n", " 'float2_e1m0fn',\n", " 'float2_e1m1fnu',\n", " 'float2_e2m0fnu',\n", " 'float32',\n", " 'float3_e1m1fn',\n", " 'float3_e1m2fnu',\n", " 'float3_e2m0fn',\n", " 'float3_e2m1fnu',\n", " 'float3_e3m0fnu',\n", " 'float4_e1m2fn',\n", " 'float4_e1m3fnu',\n", " 'float4_e2m1fn',\n", " 'float4_e2m2fnu',\n", " 'float4_e3m0fn',\n", " 'float4_e3m1fnu',\n", " 'float4_e4m0fnu',\n", " 'float5_e1m3fn',\n", " 'float5_e1m4fnu',\n", " 'float5_e2m2fn',\n", " 'float5_e2m3fnu',\n", " 'float5_e3m1fn',\n", " 'float5_e3m2fnu',\n", " 'float5_e4m0fn',\n", " 'float5_e4m1fnu',\n", " 'float5_e5m0fnu',\n", " 'float6_e1m4fn',\n", " 'float6_e1m5fnu',\n", " 'float6_e2m3fn',\n", " 'float6_e2m4fnu',\n", " 'float6_e3m2fn',\n", " 'float6_e3m3fnu',\n", " 'float6_e4m1fn',\n", " 'float6_e4m2fnu',\n", " 'float6_e5m0fn',\n", " 'float6_e5m1fnu',\n", " 'float7_e1m5fn',\n", " 'float7_e1m6fnu',\n", " 'float7_e2m4fn',\n", " 'float7_e2m5fnu',\n", " 'float7_e3m3fn',\n", " 'float7_e3m4fnu',\n", " 'float7_e4m2fn',\n", " 'float7_e4m3fnu',\n", " 'float7_e5m1fn',\n", " 'float7_e5m2fnu',\n", " 'float8_e1m6fn',\n", " 'float8_e1m7fnu',\n", " 'float8_e2m5fn',\n", " 'float8_e2m6fnu',\n", " 'float8_e3m4fn',\n", " 'float8_e3m5fnu',\n", " 'float8_e4m3fn',\n", " 'float8_e4m3fn_sdnq',\n", " 'float8_e4m3fnuz',\n", " 'float8_e4m4fnu',\n", " 'float8_e5m2',\n", " 'float8_e5m2fn',\n", " 'float8_e5m2fnuz',\n", " 'float8_e5m3fnu',\n", " 'float8_e8m0fnu',\n", " 'float9_e1m7fn',\n", " 'float9_e1m8fnu',\n", " 'float9_e2m6fn',\n", " 'float9_e2m7fnu',\n", " 'float9_e3m5fn',\n", " 'float9_e3m6fnu',\n", " 'float9_e4m4fn',\n", " 'float9_e4m5fnu',\n", " 'float9_e5m3fn',\n", " 'float9_e5m4fnu',\n", " 'fp1',\n", " 'fp10',\n", " 'fp11',\n", " 'fp12',\n", " 'fp13',\n", " 'fp14',\n", " 'fp15',\n", " 'fp16',\n", " 'fp2',\n", " 'fp3',\n", " 'fp32',\n", " 'fp4',\n", " 'fp5',\n", " 'fp6',\n", " 'fp7',\n", " 'fp8',\n", " 'fp9',\n", " 'int1',\n", " 'int10',\n", " 'int11',\n", " 'int12',\n", " 'int13',\n", " 'int14',\n", " 'int15',\n", " 'int16',\n", " 'int2',\n", " 'int3',\n", " 'int32',\n", " 'int4',\n", " 'int5',\n", " 'int6',\n", " 'int7',\n", " 'int8',\n", " 'int9',\n", " 'ufp1',\n", " 'ufp10',\n", " 'ufp11',\n", " 'ufp12',\n", " 'ufp13',\n", " 'ufp14',\n", " 'ufp15',\n", " 'ufp16',\n", " 'ufp2',\n", " 'ufp3',\n", " 'ufp4',\n", " 'ufp5',\n", " 'ufp6',\n", " 'ufp7',\n", " 'ufp8',\n", " 'ufp9',\n", " 'uint1',\n", " 'uint10',\n", " 'uint11',\n", " 'uint12',\n", " 'uint13',\n", " 'uint14',\n", " 'uint15',\n", " 'uint16',\n", " 'uint2',\n", " 'uint3',\n", " 'uint32',\n", " 'uint4',\n", " 'uint5',\n", " 'uint6',\n", " 'uint7',\n", " 'uint8',\n", " 'uint9'}" ] }, "metadata": {}, "execution_count": 7 } ] }, { "cell_type": "code", "source": [ "pipe" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "iJaa_thrWa22", "outputId": "f8109023-9b9e-4b84-ba45-c518a1b36579" }, "execution_count": null, "outputs": [ { "output_type": "execute_result", "data": { "text/plain": [ "Flux2KleinPipeline {\n", " \"_class_name\": \"Flux2KleinPipeline\",\n", " \"_diffusers_version\": \"0.38.0\",\n", " \"_name_or_path\": \"black-forest-labs/FLUX.2-klein-9B\",\n", " \"is_distilled\": true,\n", " \"scheduler\": [\n", " \"diffusers\",\n", " \"FlowMatchEulerDiscreteScheduler\"\n", " ],\n", " \"text_encoder\": [\n", " \"transformers\",\n", " \"Qwen3ForCausalLM\"\n", " ],\n", " \"tokenizer\": [\n", " \"transformers\",\n", " \"Qwen2Tokenizer\"\n", " ],\n", " \"transformer\": [\n", " \"diffusers\",\n", " \"Flux2Transformer2DModel\"\n", " ],\n", " \"vae\": [\n", " \"diffusers\",\n", " \"AutoencoderKLFlux2\"\n", " ]\n", "}" ] }, "metadata": {}, "execution_count": 18 } ] }, { "cell_type": "code", "source": [ "import torch, gc\n", "from sdnq import SDNQConfig , sdnq_post_load_quant\n", "\n", "# Apply SDNQ to transformer\n", "pipe.transformer = sdnq_post_load_quant(\n", " pipe.transformer,\n", " use_dynamic_quantization=True,\n", " weights_dtype=\"uint2\",\n", " dynamic_loss_threshold=1e-1,\n", " use_svd=True,\n", " group_size=0,\n", " svd_steps=16,\n", " quantization_device=\"cuda\",\n", " return_device=\"cpu\",\n", " quant_conv=False,\n", " quant_embedding=False,\n", ")\n", "\n", "import torch , gc\n", "with torch.no_grad():\n", " transformer_save_path = \"/content/transformer\"\n", " pipe.transformer.save_pretrained(transformer_save_path,safe_serialization=True)\n", " print(f\"✅ pipe.transformer saved to: {transformer_save_path}\")\n", "#----#\n", "gc.collect()\n", "torch.cuda.empty_cache()" ], "metadata": { "id": "hWAkx7u3XfdY" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "import gc , torch\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "import torch, gc\n", "from sdnq import SDNQConfig , sdnq_post_load_quant\n", "\n", "# Apply SDNQ to transformer\n", "pipe.text_encoder = sdnq_post_load_quant(\n", " pipe.text_encoder,\n", " use_dynamic_quantization=True,\n", " weights_dtype=\"uint2\",\n", " dynamic_loss_threshold=1e-1,\n", " use_svd=False,\n", " group_size=0,\n", " quantization_device=\"cuda\",\n", " return_device=\"cpu\",\n", " quant_conv=False,\n", " quant_embedding=False,\n", ")\n", "\n", "import torch\n", "with torch.no_grad():\n", " text_encoder_save_path = \"/content/text_encoder\"\n", " pipe.text_encoder.save_pretrained(text_encoder_save_path,safe_serialization=True)\n", " print(f\"✅ pipe.text_encoder saved to: {text_encoder_save_path}\")\n", "#----#\n", "gc.collect()\n", "torch.cuda.empty_cache()" ], "metadata": { "id": "uSPLPAVwgXp_" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "import gc , torch\n", "gc.collect()\n", "torch.cuda.empty_cache()" ], "metadata": { "id": "Xvnz0s2AV0VW" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "import torch, gc\n", "from sdnq import SDNQConfig , sdnq_post_load_quant\n", "\n", "import torch\n", "with torch.no_grad():\n", " vae_save_path = \"/content/vae\"\n", " pipe.vae.save_pretrained(vae_save_path,safe_serialization=True)\n", " print(f\"✅ pipe.vae saved to: {vae_save_path}\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "4PBgpLyDlwn4", "outputId": "eab87add-9978-45db-e585-ced9498ee6c8" }, "execution_count": 9, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "✅ pipe.vae saved to: /content/vae\n" ] } ] }, { "cell_type": "code", "source": [ "import torch\n", "with torch.no_grad():\n", " scheduler_save_path = \"/content/scheduler\"\n", " pipe.scheduler.save_pretrained(scheduler_save_path,safe_serialization=True)\n", " print(f\"✅ pipe.scheduler saved to: {scheduler_save_path}\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "nxJZeqwlmgcR", "outputId": "1bfff9be-6236-4663-ff8a-ffb4394be809" }, "execution_count": 10, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "✅ pipe.scheduler saved to: /content/scheduler\n" ] } ] }, { "cell_type": "code", "source": [ "import torch\n", "with torch.no_grad():\n", " tokenizer_save_path = \"/content/tokenizer\"\n", " pipe.tokenizer.save_pretrained(tokenizer_save_path,safe_serialization=True)\n", " print(f\"✅ pipe.tokenizer saved to: {tokenizer_save_path}\")" ], "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "ZGOf4WtqFH6g", "outputId": "571daca1-dcf9-44a8-ecd7-9c460276964f" }, "execution_count": 11, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "✅ pipe.tokenizer saved to: /content/tokenizer\n" ] } ] }, { "cell_type": "code", "metadata": { "id": "fbf8ead6" }, "source": [ "import os\n", "from huggingface_hub import HfApi, login, create_repo\n", "from google.colab import userdata\n", "from huggingface_hub.utils import HfHubHTTPError\n", "\n", "# Retrieve hf_token from Colab secrets\n", "hf_token = userdata.get(\"HF_TOKEN\")\n", "if not hf_token:\n", " raise ValueError(\"HF_TOKEN not found in Google Colab secrets. Please ensure it is set.\")\n", "\n", "# Login to Hugging Face Hub\n", "print(\"\\nLogging into Hugging Face...\")\n", "login(token=hf_token)\n", "\n", "api = HfApi()\n", "\n", "# Define the target repository ID as specified by the user\n", "repo_id = \"codeShare/FLUX.2-klein-9b-SDNQ-2bit\"\n", "\n", "# Check if repository exists, if not, create it\n", "print(f\"Checking if repository {repo_id} exists...\")\n", "try:\n", " if not api.repo_exists(repo_id=repo_id, repo_type=\"model\"):\n", " print(f\"Repository {repo_id} not found. Creating it...\")\n", " create_repo(repo_id=repo_id, repo_type=\"model\", private=False, token=hf_token)\n", " print(f\"Repository {repo_id} created successfully.\")\n", " else:\n", " print(f\"Repository {repo_id} already exists.\")\n", "except HfHubHTTPError as e:\n", " print(f\"An error occurred while checking or creating the repository: {e}\")\n", " raise # Re-raise other HTTP errors\n", "\n", "# Define paths for the folders to upload\n", "folders_to_upload = [\"/content/vae\", \"/content/tokenizer\", \"/content/scheduler\"]\n", "\n", "# Define the source path for model.safetensors.index.json and its target name in the repo\n", "model_index_source_path = \"/content/text_encoder/model.safetensors.index.json\"\n", "model_index_target_filename = \"model_index.json\"\n", "\n", "print(f\"\\nUploading specified components to {repo_id}...\")\n", "\n", "# Upload each specified folder\n", "for folder_path in folders_to_upload:\n", " if os.path.isdir(folder_path):\n", " print(f\"Uploading folder: {folder_path}...\")\n", " api.upload_folder(\n", " folder_path=folder_path,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " commit_message=f\"Upload {os.path.basename(folder_path)} component\",\n", " )\n", " print(f\"✅ Folder {folder_path} uploaded.\")\n", " else:\n", " print(f\"⚠️ Folder not found, skipping: {folder_path}\")\n", "\n", "# Upload model.safetensors.index.json as model_index.json\n", "if os.path.exists(model_index_source_path):\n", " print(f\"Uploading file: {model_index_source_path} as {model_index_target_filename}...\")\n", " api.upload_file(\n", " path_or_fileobj=model_index_source_path,\n", " path_in_repo=model_index_target_filename,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " commit_message=f\"Upload {model_index_target_filename}\",\n", " )\n", " print(f\"✅ File {model_index_source_path} uploaded as {model_index_target_filename}.\")\n", "else:\n", " print(f\"⚠️ {model_index_source_path} not found, skipping upload of {model_index_target_filename}.\")\n", "\n", "print(\"\\n✅ All specified components processed for upload to Hugging Face Hub.\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "533c64b4" }, "source": [ "import torch\n", "import gc\n", "import os\n", "import shutil\n", "from huggingface_hub import login\n", "from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel, AutoencoderKLFlux2\n", "from transformers import Qwen2Tokenizer, Qwen3ForCausalLM\n", "from sdnq import SDNQConfig, sdnq_post_load_quant # Ensure sdnq is imported\n", "\n", "# Define paths for local saved components\n", "text_encoder_local_path = \"/content/text_encoder\"\n", "transformer_local_path = \"/content/transformer\"\n", "vae_local_path = \"/content/vae\"\n", "\n", "# Define the temporary directory to assemble the pipeline\n", "rebuilt_pipeline_dir = \"/content/rebuilt_pipeline_temp\"\n", "\n", "print(f\"🔄 Rebuilding pipeline components in {rebuilt_pipeline_dir}...\")\n", "\n", "# Clean up and create the temporary directory\n", "shutil.rmtree(rebuilt_pipeline_dir, ignore_errors=True)\n", "os.makedirs(rebuilt_pipeline_dir, exist_ok=True)\n", "\n", "# Step 1: Load and save tokenizer and scheduler from the original model\n", "# These components were not SDNQ'd or individually saved as models, so we fetch them from the source.\n", "print(\"Loading original pipeline briefly to extract tokenizer and scheduler...\")\n", "original_pipe_for_components = Flux2KleinPipeline.from_pretrained(\n", " \"black-forest-labs/FLUX.2-klein-9B\",\n", " torch_dtype=torch.bfloat16, # Consistent with initial load\n", " low_cpu_mem_usage=True,\n", " device_map=\"cpu\",\n", ")\n", "\n", "original_pipe_for_components.tokenizer.save_pretrained(os.path.join(rebuilt_pipeline_dir, \"tokenizer\"))\n", "original_pipe_for_components.scheduler.save_pretrained(os.path.join(rebuilt_pipeline_dir, \"scheduler\"))\n", "\n", "del original_pipe_for_components\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "\n", "# Step 2: Copy the locally SDNQ-saved model components into the temporary pipeline directory\n", "print(f\"Copying SDNQ-applied text_encoder, transformer, vae to {rebuilt_pipeline_dir}...\")\n", "\n", "# Ensure the subdirectories exist in the rebuilt_pipeline_dir before copying\n", "os.makedirs(os.path.join(rebuilt_pipeline_dir, \"text_encoder\"), exist_ok=True)\n", "os.makedirs(os.path.join(rebuilt_pipeline_dir, \"transformer\"), exist_ok=True)\n", "os.makedirs(os.path.join(rebuilt_pipeline_dir, \"vae\"), exist_ok=True)\n", "\n", "# Copy contents of text_encoder_local_path into rebuilt_pipeline_dir/text_encoder\n", "for item_name in os.listdir(text_encoder_local_path):\n", " s = os.path.join(text_encoder_local_path, item_name)\n", " d = os.path.join(rebuilt_pipeline_dir, \"text_encoder\", item_name)\n", " if os.path.isdir(s):\n", " shutil.copytree(s, d, dirs_exist_ok=True)\n", " else:\n", " shutil.copy2(s, d)\n", "\n", "# Copy contents of transformer_local_path into rebuilt_pipeline_dir/transformer\n", "for item_name in os.listdir(transformer_local_path):\n", " s = os.path.join(transformer_local_path, item_name)\n", " d = os.path.join(rebuilt_pipeline_dir, \"transformer\", item_name)\n", " if os.path.isdir(s):\n", " shutil.copytree(s, d, dirs_exist_ok=True)\n", " else:\n", " shutil.copy2(s, d)\n", "\n", "# Copy contents of vae_local_path into rebuilt_pipeline_dir/vae\n", "for item_name in os.listdir(vae_local_path):\n", " s = os.path.join(vae_local_path, item_name)\n", " d = os.path.join(rebuilt_pipeline_dir, \"vae\", item_name)\n", " if os.path.isdir(s):\n", " shutil.copytree(s, d, dirs_exist_ok=True)\n", " else:\n", " shutil.copy2(s, d)\n", "\n", "# Step 3: Load the full pipeline from the temporary directory\n", "print(f\"Loading full pipeline from {rebuilt_pipeline_dir}...\")\n", "new_pipe = Flux2KleinPipeline.from_pretrained(\n", " rebuilt_pipeline_dir,\n", " torch_dtype=torch.bfloat16, # Consistent with initial load\n", " low_cpu_mem_usage=True,\n", " device_map=\"cpu\", # Important for offloading\n", ")\n", "print(\"✅ Base pipeline rebuilt from local components.\")\n", "\n", "# Step 4: Re-apply SDNQ to the loaded components\n", "# When loaded via from_pretrained, the base models are loaded, not the SDNQ wrappers.\n", "print(\"🔥 Re-applying SDNQ optimizations to text_encoder, transformer, and vae...\")\n", "sdnq_params = dict(\n", " use_dynamic_quantization=True,\n", " weights_dtype=\"uint4\",\n", " dynamic_loss_threshold=1e-2,\n", " use_svd=True, # Set to True for consistency with original SDNQ application\n", " group_size=0,\n", " quantization_device=\"cuda\", # Set to cuda for consistency with original SDNQ application\n", " return_device=\"cpu\",\n", " quant_conv=False,\n", " quant_embedding=False,\n", ")\n", "\n", "new_pipe.transformer = sdnq_post_load_quant(new_pipe.transformer, **sdnq_params)\n", "new_pipe.text_encoder = sdnq_post_load_quant(new_pipe.text_encoder, **sdnq_params)\n", "new_pipe.vae = sdnq_post_load_quant(new_pipe.vae, **sdnq_params)\n", "print(\"✅ SDNQ re-applied to all necessary components.\")\n", "\n", "gc.collect()\n", "torch.cuda.empty_cache()\n", "\n", "# Step 5: Login to Hugging Face and push the new pipeline\n", "print(\"\\nLogging into Hugging Face...\")\n", "login(token=userdata.get(\"HF_TOKEN\")) # hf_token is available from CELL 1\n", "\n", "print(\"\\nPushing rebuilt and SDNQ-applied pipeline to Hugging Face Hub...\")\n", "new_pipe.push_to_hub(\n", " repo_id=\"codeShare/FLUX.2-klein-9b-SDNQ-2bit\", # Using the 2bit repo_id as per request\n", " safe_serialization=True,\n", " commit_message=\"Rebuilt and pushed SDNQ Flux2 Klein 9b (2bit) after local save/crash\"\n", ")\n", "\n", "print(\"\\n✅ Rebuilt pipeline pushed to Hugging Face Hub.\")\n", "\n", "# Update the global pipe object to the newly built one\n", "pipe = new_pipe" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "3bdfe2e2" }, "source": [ "import os\n", "from huggingface_hub import HfApi, login, create_repo\n", "from google.colab import userdata\n", "from huggingface_hub.utils import HfHubHTTPError\n", "\n", "# Retrieve hf_token from Colab secrets\n", "hf_token = userdata.get(\"HF_TOKEN\")\n", "if not hf_token:\n", " raise ValueError(\"HF_TOKEN not found in Google Colab secrets. Please ensure it is set.\")\n", "\n", "# Login to Hugging Face Hub\n", "print(\"\\nLogging into Hugging Face...\")\n", "login(token=hf_token)\n", "\n", "api = HfApi()\n", "\n", "# Define the target repository ID\n", "repo_id = \"codeShare/FLUX.2-klein-9b-SDNQ-2bit\" # Or \"codeShare/FLUX.2-klein-9b-SDNQ-4bit\" if that was the intended repo\n", "\n", "# Check if repository exists, if not, create it\n", "print(f\"Checking if repository {repo_id} exists...\")\n", "try:\n", " if not api.repo_exists(repo_id=repo_id, repo_type=\"model\"):\n", " print(f\"Repository {repo_id} not found. Creating it...\")\n", " create_repo(repo_id=repo_id, repo_type=\"model\", private=False, token=hf_token)\n", " print(f\"Repository {repo_id} created successfully.\")\n", " else:\n", " print(f\"Repository {repo_id} already exists.\")\n", "except HfHubHTTPError as e:\n", " print(f\"An error occurred while checking or creating the repository: {e}\")\n", " raise # Re-raise other HTTP errors\n", "\n", "# Define local paths for the folders to upload\n", "local_folders_to_upload = [\"/content/vae\", \"/content/tokenizer\", \"/content/scheduler\"]\n", "\n", "# Define the source path for model.safetensors.index.json and its target name in the repo\n", "model_index_source_path = \"/content/text_encoder/model.safetensors.index.json\" # Assuming text_encoder still exists\n", "model_index_target_filename = \"model_index.json\"\n", "\n", "print(f\"\\nUploading specified components to {repo_id}...\")\n", "\n", "# Upload each specified folder to a corresponding path in the repo\n", "for local_folder_path in local_folders_to_upload:\n", " folder_name = os.path.basename(local_folder_path)\n", " if os.path.isdir(local_folder_path):\n", " print(f\"Uploading folder: {local_folder_path} to repo path: {folder_name}...\")\n", " api.upload_folder(\n", " folder_path=local_folder_path,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " path_in_repo=folder_name, # Uploads contents of local_folder_path into a folder named folder_name in the repo\n", " commit_message=f\"Upload {folder_name} component\",\n", " )\n", " print(f\"✅ Folder {local_folder_path} uploaded to {repo_id}/{folder_name}.\")\n", " else:\n", " print(f\"⚠️ Local folder not found, skipping: {local_folder_path}\")\n", "\n", "# Upload model.safetensors.index.json as model_index.json to the root of the repo\n", "if os.path.exists(model_index_source_path):\n", " print(f\"Uploading file: {model_index_source_path} as {model_index_target_filename} to the repo root...\")\n", " api.upload_file(\n", " path_or_fileobj=model_index_source_path,\n", " path_in_repo=model_index_target_filename,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " commit_message=f\"Upload {model_index_target_filename}\",\n", " )\n", " print(f\"✅ File {model_index_source_path} uploaded.\")\n", "else:\n", " print(f\"⚠️ {model_index_source_path} not found, skipping upload of {model_index_target_filename}.\")\n", "\n", "print(\"\\n✅ All specified components processed for upload to Hugging Face Hub.\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "a36d8446" }, "source": [ "import os\n", "from huggingface_hub import HfApi, login, create_repo\n", "from google.colab import userdata\n", "from huggingface_hub.utils import HfHubHTTPError\n", "\n", "# Retrieve hf_token from Colab secrets\n", "hf_token = userdata.get(\"HF_TOKEN\")\n", "if not hf_token:\n", " raise ValueError(\"HF_TOKEN not found in Google Colab secrets. Please ensure it is set.\")\n", "\n", "# Login to Hugging Face Hub\n", "print(\"\\nLogging into Hugging Face...\")\n", "login(token=hf_token)\n", "\n", "api = HfApi()\n", "\n", "# Define the target repository ID\n", "repo_id = \"codeShare/FLUX.2-klein-9b-SDNQ-2bit\"\n", "# Check if repository exists, if not, create it\n", "print(f\"Checking if repository {repo_id} exists...\")\n", "try:\n", " if not api.repo_exists(repo_id=repo_id, repo_type=\"model\"):\n", " print(f\"Repository {repo_id} not found. Creating it...\")\n", " create_repo(repo_id=repo_id, repo_type=\"model\", private=False, token=hf_token)\n", " print(f\"Repository {repo_id} created successfully.\")\n", " else:\n", " print(f\"Repository {repo_id} already exists.\")\n", "except HfHubHTTPError as e:\n", " print(f\"An error occurred while checking or creating the repository: {e}\")\n", " raise # Re-raise other HTTP errors\n", "\n", "# Define local paths for the folders to upload\n", "local_folders_to_upload = [\"/content/vae\", \"/content/tokenizer\", \"/content/scheduler\"]\n", "\n", "# Define the source path for model.safetensors.index.json and its target name in the repo\n", "model_index_source_path = \"/content/text_encoder/model.safetensors.index.json\" # Assuming text_encoder still exists\n", "model_index_target_filename = \"model_index.json\"\n", "\n", "print(f\"\\nUploading specified components to {repo_id}...\")\n", "\n", "# Upload each specified folder to a corresponding path in the repo\n", "for local_folder_path in local_folders_to_upload:\n", " folder_name = os.path.basename(local_folder_path)\n", " if os.path.isdir(local_folder_path):\n", " print(f\"Uploading folder: {local_folder_path} to repo path: {folder_name}...\")\n", " api.upload_folder(\n", " folder_path=local_folder_path,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " path_in_repo=folder_name, # Uploads contents of local_folder_path into a folder named folder_name in the repo\n", " commit_message=f\"Upload {folder_name} component\",\n", " )\n", " print(f\"✅ Folder {local_folder_path} uploaded to {repo_id}/{folder_name}.\")\n", " else:\n", " print(f\"⚠️ Local folder not found, skipping: {local_folder_path}\")\n", "\n", "# Upload model.safetensors.index.json as model_index.json to the root of the repo\n", "if os.path.exists(model_index_source_path):\n", " print(f\"Uploading file: {model_index_source_path} as {model_index_target_filename} to the repo root...\")\n", " api.upload_file(\n", " path_or_fileobj=model_index_source_path,\n", " path_in_repo=model_index_target_filename,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " commit_message=f\"Upload {model_index_target_filename}\",\n", " )\n", " print(f\"✅ File {model_index_source_path} uploaded as {repo_id}/{model_index_target_filename}.\")\n", "else:\n", " print(f\"⚠️ {model_index_source_path} not found, skipping upload of {model_index_target_filename}.\")\n", "\n", "print(\"\\n✅ All specified components processed for upload to Hugging Face Hub.\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "f58c3a41" }, "source": [ "import os\n", "from huggingface_hub import HfApi, login\n", "from google.colab import userdata\n", "\n", "hf_token = userdata.get(\"HF_TOKEN\")\n", "login(token=hf_token)\n", "\n", "api = HfApi()\n", "\n", "repo_id = \"codeShare/FLUX.2-klein-9b-SDNQ-2bit\"\n", "source_folder_path = \"/content/transformer\"\n", "target_folder_in_repo = \"transformer\"\n", "\n", "print(f\"\\nUploading folder: {source_folder_path} to repo path: {target_folder_in_repo} in {repo_id}...\")\n", "\n", "if os.path.isdir(source_folder_path):\n", " api.upload_folder(\n", " folder_path=source_folder_path,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " path_in_repo=target_folder_in_repo,\n", " commit_message=f\"Upload {target_folder_in_repo} component\"\n", " )\n", " print(f\"✅ Folder {source_folder_path} uploaded to {repo_id}/{target_folder_in_repo}.\")\n", "else:\n", " print(f\"⚠️ Source folder not found, skipping: {source_folder_path}\")\n", "\n", "print(\"\\n✅ Transformer component upload process complete.\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "ae55ccc9" }, "source": [ "import os\n", "from huggingface_hub import HfApi, login\n", "from google.colab import userdata\n", "\n", "hf_token = userdata.get(\"HF_TOKEN\")\n", "login(token=hf_token)\n", "\n", "api = HfApi()\n", "\n", "repo_id = \"codeShare/FLUX.2-klein-9b-SDNQ-2bit\"\n", "source_folder_path = \"/content/text_encoder\"\n", "target_folder_in_repo = \"text_encoder\"\n", "\n", "print(f\"\\nUploading folder: {source_folder_path} to repo path: {target_folder_in_repo} in {repo_id}...\")\n", "\n", "if os.path.isdir(source_folder_path):\n", " api.upload_folder(\n", " folder_path=source_folder_path,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " path_in_repo=target_folder_in_repo,\n", " commit_message=f\"Upload {target_folder_in_repo} component\"\n", " )\n", " print(f\"✅ Folder {source_folder_path} uploaded to {repo_id}/{target_folder_in_repo}.\")\n", "else:\n", " print(f\"⚠️ Source folder not found, skipping: {source_folder_path}\")\n", "\n", "print(\"\\n✅ Text Encoder component upload process complete.\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "10131a92" }, "source": [ "import os\n", "from huggingface_hub import HfApi, login\n", "from google.colab import userdata\n", "\n", "hf_token = userdata.get(\"HF_TOKEN\")\n", "login(token=hf_token)\n", "\n", "api = HfApi()\n", "\n", "repo_id = \"codeShare/FLUX.2-klein-9b-SDNQ-2bit\"\n", "\n", "folders_to_upload = {\n", " \"/content/vae\": \"vae\",\n", " \"/content/scheduler\": \"scheduler\",\n", " \"/content/tokenizer\": \"tokenizer\",\n", "}\n", "\n", "print(f\"\\nStarting upload of vae, scheduler, and tokenizer components to {repo_id}...\")\n", "\n", "for source_folder, target_folder_in_repo in folders_to_upload.items():\n", " print(f\"Uploading folder: {source_folder} to repo path: {target_folder_in_repo} in {repo_id}...\")\n", " if os.path.isdir(source_folder):\n", " api.upload_folder(\n", " folder_path=source_folder,\n", " repo_id=repo_id,\n", " repo_type=\"model\",\n", " path_in_repo=target_folder_in_repo,\n", " commit_message=f\"Upload {target_folder_in_repo} component\"\n", " )\n", " print(f\"✅ Folder {source_folder} uploaded to {repo_id}/{target_folder_in_repo}.\")\n", " else:\n", " print(f\"⚠️ Source folder not found, skipping: {source_folder}\")\n", "\n", "print(\"\\n✅ All requested components upload process complete.\")" ], "execution_count": null, "outputs": [] } ], "metadata": { "colab": { "provenance": [], "gpuType": "T4" }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "nbformat": 4, "nbformat_minor": 0 }