{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dca3239c-17d6-4284-a2cf-83237a55a7df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "add6682df5934d75bee02c59c4507a43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WaifuPipeline {\n",
      "  \"_class_name\": \"WaifuPipeline\",\n",
      "  \"_diffusers_version\": \"0.32.0.dev0\",\n",
      "  \"_name_or_path\": \"/home/recoilme/models/waifu-2b\",\n",
      "  \"scheduler\": [\n",
      "    \"diffusers\",\n",
      "    \"FlowMatchEulerDiscreteScheduler\"\n",
      "  ],\n",
      "  \"text_encoder\": [\n",
      "    \"transformers\",\n",
      "    \"XLMRobertaModel\"\n",
      "  ],\n",
      "  \"tokenizer\": [\n",
      "    \"transformers\",\n",
      "    \"XLMRobertaTokenizerFast\"\n",
      "  ],\n",
      "  \"transformer\": [\n",
      "    \"diffusers\",\n",
      "    \"SanaTransformer2DModel\"\n",
      "  ],\n",
      "  \"vae\": [\n",
      "    \"diffusers\",\n",
      "    \"AutoencoderKL\"\n",
      "  ]\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from diffusers import DiffusionPipeline\n",
    "from pipeline_waifu import WaifuPipeline\n",
    "\n",
    "pipe_id = \"/home/recoilme/models/waifu-2b\"\n",
    "variant = \"fp16\"\n",
    "pipe = WaifuPipeline.from_pretrained(\n",
    "    pipe_id, \n",
    "    variant=variant,\n",
    "    trust_remote_code=True\n",
    ")\n",
    "print(pipe)\n",
    "#pipe_sd.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6ebc579-0eb2-4828-89d5-b40f6d5e758e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def txt_embeds(prompt):\n",
    "    max_length = 512\n",
    "    select_index = [0] + list(range(-max_length + 1, 0))\n",
    "    \n",
    "    text_inputs = tokenizer(\n",
    "        prompt,\n",
    "        padding=\"max_length\",\n",
    "        max_length=max_length,\n",
    "        truncation=True,\n",
    "        add_special_tokens=True,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "    text_input_ids = text_inputs.input_ids\n",
    "    prompt_attention_mask = text_inputs.attention_mask\n",
    "    device = \"cuda\"\n",
    "    prompt_attention_mask = prompt_attention_mask.to(device)\n",
    "    \n",
    "    prompt_embeds = text_encoder(input_ids=text_input_ids.to(device), attention_mask=prompt_attention_mask)[0][:, select_index]\n",
    "    #print(prompt_embeds.shape)\n",
    "            #prompt_embeds = prompt_embeds[0][:, select_index]\n",
    "    prompt_attention_mask = prompt_attention_mask[:, select_index]\n",
    "    #print(prompt_attention_mask.shape)\n",
    "    return prompt_embeds, prompt_attention_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3704e4a3-6855-4a32-a682-7ec307a668a0",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'tokenizer' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m prompt \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mA young girl with long brown hair and brown eyes is standing against a brick wall, wearing a white shirt with a black sailor collar and a black pleated skirt. She is smiling and looking directly at the viewer, with sunlight filtering through the green vines behind her\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#prompt = 'нарядная новогодняя елка, красивые игрушки, звезда сверху, огоньки, на тёмном фоне' \u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m prompt_embeds, prompt_attention_mask \u001b[38;5;241m=\u001b[39m \u001b[43mtxt_embeds\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprompt\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      5\u001b[0m negative_prompt_embeds, negative_prompt_attention_mask\u001b[38;5;241m=\u001b[39mtxt_embeds(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28mprint\u001b[39m(prompt_embeds\u001b[38;5;241m.\u001b[39mshape)\n",
      "Cell \u001b[0;32mIn[3], line 5\u001b[0m, in \u001b[0;36mtxt_embeds\u001b[0;34m(prompt)\u001b[0m\n\u001b[1;32m      2\u001b[0m max_length \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m512\u001b[39m\n\u001b[1;32m      3\u001b[0m select_index \u001b[38;5;241m=\u001b[39m [\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mlist\u001b[39m(\u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m-\u001b[39mmax_length \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m0\u001b[39m))\n\u001b[0;32m----> 5\u001b[0m text_inputs \u001b[38;5;241m=\u001b[39m \u001b[43mtokenizer\u001b[49m(\n\u001b[1;32m      6\u001b[0m     prompt,\n\u001b[1;32m      7\u001b[0m     padding\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmax_length\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m      8\u001b[0m     max_length\u001b[38;5;241m=\u001b[39mmax_length,\n\u001b[1;32m      9\u001b[0m     truncation\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     10\u001b[0m     add_special_tokens\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     11\u001b[0m     return_tensors\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpt\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     12\u001b[0m )\n\u001b[1;32m     13\u001b[0m text_input_ids \u001b[38;5;241m=\u001b[39m text_inputs\u001b[38;5;241m.\u001b[39minput_ids\n\u001b[1;32m     14\u001b[0m prompt_attention_mask \u001b[38;5;241m=\u001b[39m text_inputs\u001b[38;5;241m.\u001b[39mattention_mask\n",
      "\u001b[0;31mNameError\u001b[0m: name 'tokenizer' is not defined"
     ]
    }
   ],
   "source": [
    "prompt = '1girl, solo, animal ears, bow, teeth, jacket, tail, open mouth, brown hair, orange background, bowtie, orange nails, simple background, cat ears, orange eyes, blue bow, animal ear fluff, cat tail, looking at viewer, upper body, shirt, school uniform, hood, striped bow, striped, white shirt, black jacket, blue bowtie, fingernails, long sleeves, cat girl, bangs, fangs, collared shirt, striped bowtie, short hair, tongue, hoodie, sharp teeth, facial mark, claw pose'\n",
    "prompt = 'A young girl with long brown hair and brown eyes is standing against a brick wall, wearing a white shirt with a black sailor collar and a black pleated skirt. She is smiling and looking directly at the viewer, with sunlight filtering through the green vines behind her'\n",
    "#prompt = 'нарядная новогодняя елка, красивые игрушки, звезда сверху, огоньки, на тёмном фоне' \n",
    "prompt_embeds, prompt_attention_mask = txt_embeds(prompt)\n",
    "negative_prompt_embeds, negative_prompt_attention_mask=txt_embeds(\"\")\n",
    "print(prompt_embeds.shape)\n",
    "\n",
    "image = pipe(\n",
    "    prompt = None,\n",
    "    negative_prompt = None,\n",
    "    use_resolution_binning = False,\n",
    "    prompt_embeds=prompt_embeds,\n",
    "    prompt_attention_mask= prompt_attention_mask,\n",
    "    negative_prompt_embeds = negative_prompt_embeds,\n",
    "    negative_prompt_attention_mask = negative_prompt_attention_mask,\n",
    "    height=768,\n",
    "    width=768,\n",
    "    guidance_scale=4.5,\n",
    "    num_inference_steps=24,\n",
    "    generator=torch.Generator(device=\"cuda\").manual_seed(42),\n",
    ")[0]\n",
    "\n",
    "for img in image:\n",
    "    img.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b82a561-b93b-4261-9bc7-fe867168bdd2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}