{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7e8f9dc5-d07a-4538-bc03-8953412a72fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Keyword arguments {'safety_checker': <__main__.DummyCosmosSafetyChecker object at 0x7f7e8c3fb620>} are not expected by SdxsPipeline and will be ignored.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99e2522e93064308b5dd34923a133c39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "def23cc245b2470a9012f70d5e4c78ed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/195 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'final_sigmas_type': 'sigma_min', 'sigma_data': 1.0, 'sigma_max': 80.0, 'sigma_min': 0.002} were passed to FlowMatchEulerDiscreteScheduler, but are not expected and will be ignored. Please verify your scheduler_config.json configuration file.\n",
      "Sampling: 100%|██████████| 40/40 [00:10<00:00,  3.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Готово! Изображение сохранено как output.png\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from diffusers import Cosmos2TextToImagePipeline\n",
    "\n",
    "class DummyCosmosSafetyChecker:\n",
    "    def to(self, *args, **kwargs):\n",
    "        return self\n",
    "        \n",
    "    def eval(self):\n",
    "        return self\n",
    "\n",
    "    # Обход проверки текста\n",
    "    def check_text_safety(self, prompt, *args, **kwargs):\n",
    "        return True\n",
    "\n",
    "    # Обход проверки \"видео\" (картинки из 1 кадра)\n",
    "    def check_video_safety(self, vid, *args, **kwargs):\n",
    "        # Просто возвращаем тензор обратно без изменений\n",
    "        return vid\n",
    "\n",
    "    # На всякий случай оставляем оригинальный __call__\n",
    "    def __call__(self, images, **kwargs):\n",
    "        return images, [False] * len(images)\n",
    "\n",
    "model_id = \"/workspace/sdxs-2b\"\n",
    "\n",
    "pipe = Cosmos2TextToImagePipeline.from_pretrained(\n",
    "    model_id,\n",
    "    safety_checker=DummyCosmosSafetyChecker(), \n",
    "    torch_dtype=torch.bfloat16                 \n",
    ")\n",
    "pipe.to(\"cuda\")\n",
    "\n",
    "prompt = \"In a serene garden, two young girls stand side by side, their youthful energy palpable. The girl on the left, adorned with a blue dress and a matching blue flower in her hair, gazes directly at the viewer, her eyes sparkling with curiosity.\"#\"There is a young male character standing against a vibrant, colorful graffiti wall. he is wearing a  hat, a  jacket adorned with gold accents, and black shorts.\"\n",
    "negative_prompt = \"The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.\"\n",
    "\n",
    "# 3. Генерируем изображение\n",
    "output = pipe(\n",
    "    height = 1024,\n",
    "    width=1024,\n",
    "    prompt=prompt, \n",
    "    negative_prompt=negative_prompt, \n",
    "    generator=torch.Generator(device=\"cuda\").manual_seed(1)\n",
    ").images[0]\n",
    "\n",
    "output.save(\"output.png\")\n",
    "print(\"Готово! Изображение сохранено как output.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a173167-6c28-4bbd-8879-1375e0fd37f0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 (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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}