{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ff645d46-9401-4cd1-8dc1-d5eb98a7b1e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Load schema\n",
    "with open(\"schema.json\") as f:\n",
    "    SCHEMA = json.load(f)\n",
    "\n",
    "SYSTEM_PROMPT = f\"\"\"\n",
    "You are an SQL generator for retail analytics (AWS Athena), only refer the tables JSON schema shared below.\n",
    "\n",
    "Tables:\n",
    "{json.dumps(SCHEMA, indent=2)}\n",
    "\n",
    "Your output must strictly be ** EXACTLY ONE JSON OBJECT ** with following structure:\n",
    "{{ \n",
    "  \"sql_query\": \"...\",\n",
    "  \"query_planning\": \"...\",\n",
    "  \"explanation\": \"...\"\n",
    "}}\n",
    "\n",
    "\n",
    "Here are some Examples of User & Assisant conversations:\n",
    "\n",
    "1) Example 1:\n",
    "User: \"get all competitor products in Mobiles category from flipkart\"\n",
    "Assistant:\n",
    "{{\n",
    "  \"sql_query\": \"SELECT sku_uuid, comp_sku, brand, category, store_name, source_name, comp_price\\nFROM pdp_newdev.product_warehouse\\nWHERE category = 'Mobiles' AND lower(source_name) = 'flipkart';\",\n",
    "  \"query_planning\": [\n",
    "    \"Filter product_warehouse by category='Mobiles'\",\n",
    "    \"Filter source_name='flipkart'\",\n",
    "    \"Select SKU, brand, category, pricing\"\n",
    "  ],\n",
    "  \"explanation\": \"Lists competitor products in the Mobiles category on Flipkart.\"\n",
    "}}\n",
    "\n",
    "\n",
    "2) Example 2:\n",
    "User: \"show base products loaded on 2024-04-15 under brand Boat\"\n",
    "Assistant:\n",
    "{{\n",
    "  \"sql_query\": \"SELECT base_sku, product_title, brand, category, price, load_date\\nFROM bungee_customercatalog.athena_auroradb_catalog\\nWHERE load_date = DATE '2024-04-15' AND brand = 'Boat';\",\n",
    "  \"query_planning\": [\n",
    "    \"Filter athena_auroradb_catalog by load_date and brand='Boat'\",\n",
    "    \"Select SKU, title, category, price, load_date\"\n",
    "  ],\n",
    "  \"explanation\": \"Retrieves Boat-brand base products loaded on 2024-04-15.\"\n",
    "}}\n",
    "\n",
    "3) Example 3:\n",
    "User: \"retrieve Chewy matches with active status\"\n",
    "Assistant:\n",
    "{{\n",
    "  \"sql_query\": \"SELECT a.*\\nFROM (\\n  SELECT\\n    SPLIT_PART(comp_source_store, '_', 1) AS source,\\n    base_sku        AS chewy_sku,\\n    comp_sku        AS sku,\\n    ''              AS priority,\\n    comp_upc\\n  FROM match_library.match_library_snapshot\\n  WHERE load_date = (\\n    SELECT MAX(load_date)\\n    FROM match_library.match_library_snapshot\\n  )\\n    AND base_source_store = 'chewy_chewy'\\n    AND active = TRUE\\n    AND match_status = 'product_found'\\n  GROUP BY 1,2,3,4,5\\n) a\\nJOIN (\\n  SELECT sku\\n  FROM bungee_customercatalog.athena_auroradb_catalog\\n  WHERE source = 'chewy'\\n    AND load_date = (\\n      SELECT MAX(load_date)\\n      FROM bungee_customercatalog.athena_auroradb_catalog\\n      WHERE source = 'chewy'\\n    )\\n) b\\n  ON a.chewy_sku = b.sku;\",\n",
    "  \"query_planning\": [\n",
    "    \"Filter match_library_snapshot to the latest Chewy load_date\",\n",
    "    \"Restrict to active, product_found matches and group key fields\",\n",
    "    \"Join to latest Chewy catalog on SKU\"\n",
    "  ],\n",
    "  \"explanation\": \"Fetches active, verified Chewy product matches using the most recent snapshot.\"\n",
    "}}\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "22e393f9-a6b4-4b26-ae84-f81c865ea2f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(SYSTEM_PROMPT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e402d20e-7aff-4f9d-a0d4-3eb8b03b2030",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:__main__:Chat history cleared (UI only).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Serving Flask app '__main__'\n",
      " * Debug mode: off\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:werkzeug:\u001b[31m\u001b[1mWARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.\u001b[0m\n",
      " * Running on http://127.0.0.1:5000\n",
      "INFO:werkzeug:\u001b[33mPress CTRL+C to quit\u001b[0m\n",
      "INFO:werkzeug:127.0.0.1 - - [18/Apr/2025 08:43:48] \"GET / HTTP/1.1\" 200 -\n",
      "INFO:__main__:hi\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "import logging\n",
    "from flask import Flask, request, render_template, jsonify\n",
    "import ollama\n",
    "from pyathena import connect\n",
    "# from system_prompt import SYSTEM_PROMPT\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "OLLAMA_MODEL  = os.getenv(\"OLLAMA_MODEL\", \"mistral:latest\")\n",
    "ATHENA_DB     = os.getenv(\"ATHENA_DATABASE\")\n",
    "ATHENA_OUTPUT = os.getenv(\"ATHENA_OUTPUT_LOCATION\")\n",
    "AWS_REGION    = os.getenv(\"AWS_REGION\", \"us-east-1\")\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# UI history only (never sent wholesale to the model)\n",
    "chat_history = []\n",
    "def init_history():\n",
    "    global chat_history\n",
    "    chat_history = []\n",
    "    logger.info(\"Chat history cleared (UI only).\")\n",
    "\n",
    "init_history()\n",
    "\n",
    "def run_athena_query(sql: str) -> list:\n",
    "    conn = connect(\n",
    "        s3_staging_dir='s3://ml.dev/athena_query_op/',\n",
    "        region_name='us-east-1'\n",
    "    )\n",
    "    cursor = conn.cursor()\n",
    "    cursor.execute(sql)\n",
    "    cols = [d[0] for d in cursor.description]\n",
    "    rows = cursor.fetchall()\n",
    "    return [dict(zip(cols, row)) for row in rows]\n",
    "\n",
    "@app.route(\"/\", methods=[\"GET\"])\n",
    "def index():\n",
    "    # Render UI with system + history\n",
    "    full = [{\"role\": \"system\", \"content\": SYSTEM_PROMPT}] + chat_history\n",
    "    return render_template(\"index.html\", chat_history=full)\n",
    "\n",
    "@app.route(\"/chat\", methods=[\"POST\"])\n",
    "def chat():\n",
    "    global chat_history\n",
    "    user_text = request.form.get(\"message\", \"\").strip()\n",
    "    logger.info(str(user_text))\n",
    "    if not user_text:\n",
    "        return jsonify(error=\"No message\"), 400\n",
    "\n",
    "    # Add to UI history\n",
    "    chat_history.append({\"role\": \"user\", \"content\": user_text})\n",
    "\n",
    "    # ALWAYS send system prompt first, then the *single* user message\n",
    "    payload = [\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\",   \"content\": user_text}\n",
    "    ]\n",
    "    # logger.info(\"Payload to model: %s\", payload)\n",
    "\n",
    "    resp = ollama.chat(\n",
    "        model=OLLAMA_MODEL,\n",
    "        messages=payload,\n",
    "        options={\"temperature\": 0}\n",
    "    )\n",
    "    out = resp.message.content\n",
    "    logger.info(\"Model output: %s\", out)\n",
    "\n",
    "    # Try to parse JSON\n",
    "    try:\n",
    "        result = json.loads(out)\n",
    "        sql = result.get(\"sql_query\", \"\").strip()\n",
    "    except Exception:\n",
    "        # Fallback to raw\n",
    "        chat_history.append({\"role\": \"assistant\", \"content\": out})\n",
    "        full = [{\"role\":\"system\",\"content\":SYSTEM_PROMPT}] + chat_history\n",
    "        return jsonify(chat_history=full)\n",
    "\n",
    "    # Optionally run it\n",
    "    if sql:\n",
    "        try:\n",
    "            result[\"athena_results\"] = run_athena_query(sql)\n",
    "        except Exception as e:\n",
    "            result[\"error\"] = str(e)\n",
    "\n",
    "    formatted = json.dumps(result, indent=2)\n",
    "    chat_history.append({\"role\": \"assistant\", \"content\": formatted})\n",
    "\n",
    "    full = [{\"role\":\"system\",\"content\":SYSTEM_PROMPT}] + chat_history\n",
    "    return jsonify(chat_history=full)\n",
    "\n",
    "@app.route(\"/clear\", methods=[\"POST\"])\n",
    "def clear():\n",
    "    init_history()\n",
    "    full = [{\"role\":\"system\",\"content\":SYSTEM_PROMPT}] + chat_history\n",
    "    return jsonify(chat_history=full)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    app.run(debug=False, use_reloader=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "12aac119-800f-4306-9a3e-1660f7333a4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:__main__:Chat history cleared (UI only).\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "chat() takes 0 positional arguments but 1 was given",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-6822d404fdc9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0minit_history\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mchat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"total crawled products for petco and amazon both in 2025\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: chat() takes 0 positional arguments but 1 was given"
     ]
    }
   ],
   "source": [
    "init_history()\n",
    "chat(\"total crawled products for petco and amazon both in 2025\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2782ace6-6234-4990-a920-e15ff9dd3f7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0538594-88bf-4ae9-9510-3ad96e17c8fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "ListResponse(models=[Model(model='cogito:latest', modified_at=datetime.datetime(2025, 4, 15, 21, 39, 31, 7685, tzinfo=TzInfo(UTC)), digest='75b508ddece134205c621b91237310024ca983027fb7f724c2032d106419bd1b', size=4920747948, details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='8.0B', quantization_level='Q4_K_M')), Model(model='artifish/llama3.2-uncensored:latest', modified_at=datetime.datetime(2025, 4, 15, 20, 28, 19, 718626, tzinfo=TzInfo(UTC)), digest='c73bea26e00483b58ba59729b74bfc1647971dafa1fe66c890d845ae63e2c34e', size=2241006824, details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='3.6B', quantization_level='Q4_K_M')), Model(model='deepseek-coder-v2:latest', modified_at=datetime.datetime(2025, 4, 15, 6, 53, 32, 916369, tzinfo=TzInfo(UTC)), digest='63fb193b3a9b4322a18e8c6b250ca2e70a5ff531e962dbf95ba089b2566f2fa5', size=8905126121, details=ModelDetails(parent_model='', format='gguf', family='deepseek2', families=['deepseek2'], parameter_size='15.7B', quantization_level='Q4_0')), Model(model='deepseek-v3:671b', modified_at=datetime.datetime(2025, 4, 15, 6, 7, 47, 388953, tzinfo=TzInfo(UTC)), digest='5da0e2d4a9e0e3cc2cac1e952a6808f1eccdc75acd9cf4387e41c8ccc77a929a', size=404430200769, details=ModelDetails(parent_model='', format='gguf', family='deepseek2', families=['deepseek2'], parameter_size='671.0B', quantization_level='Q4_K_M')), Model(model='codellama:latest', modified_at=datetime.datetime(2025, 4, 15, 4, 47, 50, 580852, tzinfo=TzInfo(UTC)), digest='8fdf8f752f6e80de33e82f381aba784c025982752cd1ae9377add66449d2225f', size=3825910662, details=ModelDetails(parent_model='', format='gguf', family='llama', families=None, parameter_size='7B', quantization_level='Q4_0')), Model(model='nomic-embed-text:latest', modified_at=datetime.datetime(2025, 1, 3, 9, 49, 49, 322519, tzinfo=TzInfo(UTC)), digest='0a109f422b47e3a30ba2b10eca18548e944e8a23073ee3f3e947efcf3c45e59f', size=274302450, details=ModelDetails(parent_model='', format='gguf', family='nomic-bert', families=['nomic-bert'], parameter_size='137M', quantization_level='F16')), Model(model='llama3.2:latest', modified_at=datetime.datetime(2024, 12, 28, 13, 55, 10, 79145, tzinfo=TzInfo(UTC)), digest='a80c4f17acd55265feec403c7aef86be0c25983ab279d83f3bcd3abbcb5b8b72', size=2019393189, details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='3.2B', quantization_level='Q4_K_M')), Model(model='llama2:latest', modified_at=datetime.datetime(2024, 2, 29, 0, 21, 2, 69015, tzinfo=TzInfo(UTC)), digest='78e26419b4469263f75331927a00a0284ef6544c1975b826b15abdaef17bb962', size=3826793677, details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='7B', quantization_level='Q4_0')), Model(model='mistral:latest', modified_at=datetime.datetime(2024, 2, 26, 23, 49, 35, 566793, tzinfo=TzInfo(UTC)), digest='61e88e884507ba5e06c49b40e6226884b2a16e872382c2b44a42f2d119d804a5', size=4109865159, details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='7B', quantization_level='Q4_0'))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1c84683f-67d7-4316-88ae-a2e1680c85db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7d166f3f-7656-45ad-b196-9c340002637a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f9dc06df4fc44ff49ac3552b1ddf0db0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee8073a0e92a486db7880c6f8484b4ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0ce727d4c174defa6dbea5a9f3f93df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "OutOfMemoryError",
     "evalue": "CUDA out of memory. Tried to allocate 160.00 MiB. GPU 0 has a total capacity of 23.46 GiB of which 158.94 MiB is free. Process 83751 has 6.79 GiB memory in use. Including non-PyTorch memory, this process has 16.49 GiB memory in use. Of the allocated memory 16.29 GiB is allocated by PyTorch, and 37.88 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-4bfc7904887f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;31m# quick test\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mpipe\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDiffusionPipeline\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_pretrained\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"black-forest-labs/FLUX.1-dev\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mpipe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"cuda\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpipe\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Astronaut in a jungle, cold color palette, detailed, 8k\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimages\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"test.png\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/diffusers/pipelines/pipeline_utils.py\u001b[0m in \u001b[0;36mto\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    480\u001b[0m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    481\u001b[0m             \u001b[0;32melif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mis_loaded_in_4bit_bnb\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mis_loaded_in_8bit_bnb\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mis_group_offloaded\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 482\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    483\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    484\u001b[0m             if (\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/transformers/modeling_utils.py\u001b[0m in \u001b[0;36mto\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   3155\u001b[0m                     \u001b[0;34m\" `dtype` by passing the correct `torch_dtype` argument.\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3156\u001b[0m                 )\n\u001b[0;32m-> 3157\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3158\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3159\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mhalf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mto\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1172\u001b[0m                     \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1173\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1174\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconvert\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1175\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1176\u001b[0m     def register_full_backward_pre_hook(\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    778\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 780\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    781\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    782\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    803\u001b[0m             \u001b[0;31m# `with torch.no_grad():`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    804\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 805\u001b[0;31m                 \u001b[0mparam_applied\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    806\u001b[0m             \u001b[0mp_should_use_set_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparam_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    807\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mconvert\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m   1158\u001b[0m                         \u001b[0mmemory_format\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mconvert_to_format\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1159\u001b[0m                     )\n\u001b[0;32m-> 1160\u001b[0;31m                 return t.to(\n\u001b[0m\u001b[1;32m   1161\u001b[0m                     \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1162\u001b[0m                     \u001b[0mdtype\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_floating_point\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_complex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 160.00 MiB. GPU 0 has a total capacity of 23.46 GiB of which 158.94 MiB is free. Process 83751 has 6.79 GiB memory in use. Including non-PyTorch memory, this process has 16.49 GiB memory in use. Of the allocated memory 16.29 GiB is allocated by PyTorch, and 37.88 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)"
     ]
    }
   ],
   "source": [
    "from diffusers import DiffusionPipeline\n",
    "import torch\n",
    "\n",
    "# quick test\n",
    "pipe = DiffusionPipeline.from_pretrained(\"black-forest-labs/FLUX.1-dev\")\n",
    "pipe.to(\"cuda\")\n",
    "img = pipe(\"Astronaut in a jungle, cold color palette, detailed, 8k\").images[0]\n",
    "img.save(\"test.png\")\n",
    "print(\"✅ Generated test.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46d533a1-c05a-43d9-95e5-3a24bc196a50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In your terminal (or in a Jupyter cell prefixing with !)\n",
    "!pip install diffusers transformers accelerate safetensors\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1be04254-d787-458e-963a-e379d179cafa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffdcc284-a56c-45ab-b8bc-b88ab0609fa0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 5
}