colab_demo_t2i.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install diffusers transformers accelerate peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import StableDiffusionXLPipeline\n",
    "from scheduling_tcd import TCDScheduler "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\"\n",
    "base_model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
    "tcd_lora_id = \"h1t/TCD-SDXL-LoRA\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the diffusers pipeline for text-to-image generation.\n",
    "pipe = StableDiffusionXLPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, variant=\"fp16\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the scheduler to TCDScheduler.\n",
    "pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)\n",
    "\n",
    "# Load the TCD-LoRA weights for the model.\n",
    "pipe.load_lora_weights(tcd_lora_id)\n",
    "pipe.fuse_lora()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Beautiful woman, bubblegum pink, lemon yellow, minty blue, futuristic, high-detail, epic composition, watercolor.\"\n",
    "\n",
    "# Perform inference with the pipeline.\n",
    "image = pipe(\n",
    "    prompt=prompt,\n",
    "    num_inference_steps=4,\n",
    "    guidance_scale=0,\n",
    "    # Eta (referred to as `gamma` in the paper) is used to control the stochasticity in every step.\n",
    "    # A value of 0.3 often yields good results.\n",
    "    # We recommend using a higher eta when increasing the number of inference steps.\n",
    "    eta=0.3, \n",
    "    generator=torch.Generator(device=device).manual_seed(0),\n",
    ").images[0]"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}