{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# DiffuseCraft: Text-to-Image Generation on T4 Colab\n",
        "\n",
        "This script uses a custom Stable Diffusion model from Hugging Face for text-to-image generation, optimized for T4 GPU with low RAM usage.\n",
        "\n",
        "**Requirements**:\n",
        "- T4 GPU runtime in Colab\n",
        "- Hugging Face account and token (for gated models)\n",
        "\n",
        "**Features**:\n",
        "- Uses `diffusers` library with FP16 precision\n",
        "- Enables model CPU offloading for low RAM\n",
        "- Supports custom prompts and negative prompts\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Install required libraries\n",
        "!pip install -q diffusers==0.21.4 transformers==4.33.0 accelerate==0.22.0\n",
        "!pip install -q torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118\n",
        "!pip install -q xformers==0.0.22\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Import libraries\n",
        "import torch\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from huggingface_hub import login\n",
        "import os\n",
        "\n",
        "# Set Hugging Face token (replace with your token)\n",
        "os.environ['HUGGINGFACE_TOKEN'] = 'your_hf_token_here'\n",
        "login(os.environ['HUGGINGFACE_TOKEN'])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Initialize the pipeline with optimizations\n",
        "model_id = 'runwayml/stable-diffusion-v1-5'  # Replace with your custom HF model ID\n",
        "\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\n",
        "    model_id,\n",
        "    torch_dtype=torch.float16,\n",
        "    use_auth_token=True\n",
        ")\n",
        "\n",
        "# Enable optimizations for T4\n",
        "pipe = pipe.to('cuda')\n",
        "pipe.enable_attention_slicing()  # Reduces memory usage\n",
        "pipe.enable_model_cpu_offload()  # Offloads model to CPU when not in use\n",
        "\n",
        "# Optional: Enable xformers for faster inference\n",
        "try:\n",
        "    pipe.enable_xformers_memory_efficient_attention()\n",
        "except:\n",
        "    print('xformers not supported, proceeding without it.')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Define generation parameters\n",
        "prompt = 'A serene mountain landscape at sunset, vibrant colors, highly detailed'\n",
        "negative_prompt = 'blurry, low quality, artifacts, text, watermark'\n",
        "num_inference_steps = 30  # Lower steps for faster generation\n",
        "guidance_scale = 7.5\n",
        "\n",
        "# Generate image\n",
        "image = pipe(\n",
        "    prompt,\n",
        "    negative_prompt=negative_prompt,\n",
        "    num_inference_steps=num_inference_steps,\n",
        "    guidance_scale=guidance_scale,\n",
        "    height=512,\n",
        "    width=512\n",
        ").images[0]\n",
        "\n",
        "# Save and display image\n",
        "image.save('generated_image.png')\n",
        "image\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Notes\n",
        "- Replace `'your_hf_token_here'` with your Hugging Face token.\n",
        "- Replace `'runwayml/stable-diffusion-v1-5'` with your custom model ID from Hugging Face.\n",
        "- Adjust `prompt`, `negative_prompt`, `num_inference_steps`, and `guidance_scale` as needed.\n",
        "- The script uses FP16 and attention slicing to minimize RAM usage.\n",
        "- Model CPU offloading reduces VRAM requirements, ideal for T4 GPUs.\n"
      ]
    }
  ],
  "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.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}