{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# SDXL with Custom LoRA on T4 GPU\n",
        "\n",
        "This notebook sets up Stable Diffusion XL (SDXL) on a T4 GPU in Google Colab with Python 3.11, downloads the base model from Hugging Face, and applies a custom LoRA model from Hugging Face or Civitai. It generates images using the configured pipeline.\n",
        "\n",
        "**Prerequisites:**\n",
        "- Hugging Face account and token for gated model access (e.g., `stabilityai/stable-diffusion-xl-base-1.0`).\n",
        "- Civitai API key if downloading LoRA from Civitai.\n",
        "- Ensure Colab is set to T4 GPU (Runtime > Change runtime type > T4 GPU).\n",
        "\n",
        "**Note:** Replace placeholders (e.g., `YOUR_HF_TOKEN`, `YOUR_CIVITAI_API_KEY`) with your actual credentials."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Install dependencies for Python 3.11\n",
        "!pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu118\n",
        "!pip install diffusers==0.29.2 transformers==4.44.2 accelerate==0.33.0 safetensors==0.4.5\n",
        "!pip install requests\n",
        "\n",
        "# Verify Python version\n",
        "import sys\n",
        "print(sys.version)\n",
        "\n",
        "# Check GPU availability\n",
        "import torch\n",
        "print(f\"CUDA Available: {torch.cuda.is_available()}\")\n",
        "print(f\"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'No GPU'}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 1: Authenticate and Download Models\n",
        "Authenticate with Hugging Face to download the SDXL base model. Optionally, provide a Civitai API key for LoRA download."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "from huggingface_hub import login\n",
        "import os\n",
        "\n",
        "# Log in to Hugging Face\n",
        "HF_TOKEN = \"YOUR_HF_TOKEN\"  # Replace with your Hugging Face token\n",
        "login(HF_TOKEN)\n",
        "\n",
        "# Set Civitai API key (if downloading from Civitai)\n",
        "CIVITAI_API_KEY = \"YOUR_CIVITAI_API_KEY\"  # Replace with your Civitai API key or set to None\n",
        "os.environ[\"CIVITAI_API_KEY\"] = CIVITAI_API_KEY if CIVITAI_API_KEY else \"\"\n",
        "\n",
        "# Download SDXL base model\n",
        "from diffusers import StableDiffusionXLPipeline\n",
        "base_model = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
        "pipeline = StableDiffusionXLPipeline.from_pretrained(\n",
        "    base_model,\n",
        "    torch_dtype=torch.float16,\n",
        "    variant=\"fp16\",\n",
        "    use_safetensors=True\n",
        ").to(\"cuda\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 2: Download Custom LoRA\n",
        "Choose to download a LoRA model from Hugging Face or Civitai. Replace the URLs/IDs with your desired LoRA model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import requests\n",
        "import os\n",
        "\n",
        "def download_hf_lora(repo_id, filename, local_dir=\"./lora\"):\n",
        "    os.makedirs(local_dir, exist_ok=True)\n",
        "    local_path = os.path.join(local_dir, filename)\n",
        "    url = f\"https://huggingface.co/{repo_id}/resolve/main/{filename}\"\n",
        "    headers = {\"Authorization\": f\"Bearer {HF_TOKEN}\"}\n",
        "    response = requests.get(url, headers=headers)\n",
        "    response.raise_for_status()\n",
        "    with open(local_path, \"wb\") as f:\n",
        "        f.write(response.content)\n",
        "    return local_path\n",
        "\n",
        "def download_civitai_lora(model_id, filename, local_dir=\"./lora\"):\n",
        "    os.makedirs(local_dir, exist_ok=True)\n",
        "    local_path = os.path.join(local_dir, filename)\n",
        "    url = f\"https://civitai.com/api/download/models/{model_id}\"\n",
        "    headers = {\"Authorization\": f\"Bearer {os.environ['CIVITAI_API_KEY']}\"}\n",
        "    response = requests.get(url, headers=headers)\n",
        "    response.raise_for_status()\n",
        "    with open(local_path, \"wb\") as f:\n",
        "        f.write(response.content)\n",
        "    return local_path\n",
        "\n",
        "# Example: Download LoRA (choose one method)\n",
        "# Hugging Face LoRA (e.g., a hypothetical LoRA model)\n",
        "lora_path = download_hf_lora(\n",
        "    repo_id=\"username/sdxl-lora-model\",  # Replace with actual Hugging Face repo ID\n",
        "    filename=\"model.safetensors\"         # Replace with actual filename\n",
        ")\n",
        "\n",
        "# Civitai LoRA (uncomment to use)\n",
        "# lora_path = download_civitai_lora(\n",
        "#     model_id=\"123456\",               # Replace with Civitai model ID\n",
        "#     filename=\"lora_model.safetensors\"\n",
        "# )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 3: Load LoRA into Pipeline\n",
        "Load the custom LoRA weights into the SDXL pipeline."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Load LoRA weights\n",
        "pipeline.load_lora_weights(\n",
        "    lora_path,\n",
        "    adapter_name=\"custom_lora\"\n",
        ")\n",
        "\n",
        "# Enable LoRA\n",
        "pipeline.set_adapters([\"custom_lora\"], adapter_weights=[1.0])\n",
        "\n",
        "# Optimize for T4 GPU\n",
        "pipeline.enable_model_cpu_offload()\n",
        "pipeline.enable_vae_slicing()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 4: Generate Images\n",
        "Configure the prompt and generate images using the SDXL pipeline with LoRA."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "prompt = \"A futuristic cityscape at sunset, cyberpunk style, highly detailed, vibrant colors\"\n",
        "negative_prompt = \"blurry, low quality, artifacts\"\n",
        "\n",
        "images = pipeline(\n",
        "    prompt=prompt,\n",
        "    negative_prompt=negative_prompt,\n",
        "    num_inference_steps=30,\n",
        "    guidance_scale=7.5,\n",
        "    height=1024,\n",
        "    width=1024,\n",
        "    num_images_per_prompt=1\n",
        ").images\n",
        "\n",
        "# Save and display the image\n",
        "images[0].save(\"output.png\")\n",
        "images[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Notes\n",
        "- **T4 GPU Optimization**: The notebook uses `float16` precision, VAE slicing, and model CPU offloading to fit within T4 GPU memory constraints (16GB VRAM).\n",
        "- **LoRA Model**: Ensure the LoRA model is compatible with SDXL. Replace placeholder repo IDs or model IDs with actual values from Hugging Face or Civitai.\n",
        "- **Performance**: Adjust `num_inference_steps` and `guidance_scale` for quality vs. speed trade-offs.\n",
        "- **Storage**: Generated images are saved as `output.png` in the Colab environment. Download them manually if needed."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}