Format: application/octet-stream
1 {
2 "cells": [
3 {
4 "cell_type": "code",
5 "execution_count": 1,
6 "id": "ff645d46-9401-4cd1-8dc1-d5eb98a7b1e1",
7 "metadata": {},
8 "outputs": [],
9 "source": [
10 "import json\n",
11 "\n",
12 "# Load schema\n",
13 "with open(\"schema.json\") as f:\n",
14 " SCHEMA = json.load(f)\n",
15 "\n",
16 "SYSTEM_PROMPT = f\"\"\"\n",
17 "You are an SQL generator for retail analytics (AWS Athena), only refer the tables JSON schema shared below.\n",
18 "\n",
19 "Tables:\n",
20 "{json.dumps(SCHEMA, indent=2)}\n",
21 "\n",
22 "Your output must strictly be ** EXACTLY ONE JSON OBJECT ** with following structure:\n",
23 "{{ \n",
24 " \"sql_query\": \"...\",\n",
25 " \"query_planning\": \"...\",\n",
26 " \"explanation\": \"...\"\n",
27 "}}\n",
28 "\n",
29 "\n",
30 "Here are some Examples of User & Assisant conversations:\n",
31 "\n",
32 "1) Example 1:\n",
33 "User: \"get all competitor products in Mobiles category from flipkart\"\n",
34 "Assistant:\n",
35 "{{\n",
36 " \"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",
37 " \"query_planning\": [\n",
38 " \"Filter product_warehouse by category='Mobiles'\",\n",
39 " \"Filter source_name='flipkart'\",\n",
40 " \"Select SKU, brand, category, pricing\"\n",
41 " ],\n",
42 " \"explanation\": \"Lists competitor products in the Mobiles category on Flipkart.\"\n",
43 "}}\n",
44 "\n",
45 "\n",
46 "2) Example 2:\n",
47 "User: \"show base products loaded on 2024-04-15 under brand Boat\"\n",
48 "Assistant:\n",
49 "{{\n",
50 " \"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",
51 " \"query_planning\": [\n",
52 " \"Filter athena_auroradb_catalog by load_date and brand='Boat'\",\n",
53 " \"Select SKU, title, category, price, load_date\"\n",
54 " ],\n",
55 " \"explanation\": \"Retrieves Boat-brand base products loaded on 2024-04-15.\"\n",
56 "}}\n",
57 "\n",
58 "3) Example 3:\n",
59 "User: \"retrieve Chewy matches with active status\"\n",
60 "Assistant:\n",
61 "{{\n",
62 " \"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",
63 " \"query_planning\": [\n",
64 " \"Filter match_library_snapshot to the latest Chewy load_date\",\n",
65 " \"Restrict to active, product_found matches and group key fields\",\n",
66 " \"Join to latest Chewy catalog on SKU\"\n",
67 " ],\n",
68 " \"explanation\": \"Fetches active, verified Chewy product matches using the most recent snapshot.\"\n",
69 "}}\n",
70 "\"\"\"\n"
71 ]
72 },
73 {
74 "cell_type": "code",
75 "execution_count": 2,
76 "id": "22e393f9-a6b4-4b26-ae84-f81c865ea2f2",
77 "metadata": {},
78 "outputs": [],
79 "source": [
80 "# print(SYSTEM_PROMPT)"
81 ]
82 },
83 {
84 "cell_type": "code",
85 "execution_count": 3,
86 "id": "e402d20e-7aff-4f9d-a0d4-3eb8b03b2030",
87 "metadata": {},
88 "outputs": [
89 {
90 "name": "stderr",
91 "output_type": "stream",
92 "text": [
93 "INFO:__main__:Chat history cleared (UI only).\n"
94 ]
95 },
96 {
97 "name": "stdout",
98 "output_type": "stream",
99 "text": [
100 " * Serving Flask app '__main__'\n",
101 " * Debug mode: off\n"
102 ]
103 },
104 {
105 "name": "stderr",
106 "output_type": "stream",
107 "text": [
108 "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",
109 " * Running on http://127.0.0.1:5000\n",
110 "INFO:werkzeug:\u001b[33mPress CTRL+C to quit\u001b[0m\n",
111 "INFO:werkzeug:127.0.0.1 - - [18/Apr/2025 08:43:48] \"GET / HTTP/1.1\" 200 -\n",
112 "INFO:__main__:hi\n"
113 ]
114 }
115 ],
116 "source": [
117 "import os\n",
118 "import json\n",
119 "import logging\n",
120 "from flask import Flask, request, render_template, jsonify\n",
121 "import ollama\n",
122 "from pyathena import connect\n",
123 "# from system_prompt import SYSTEM_PROMPT\n",
124 "\n",
125 "app = Flask(__name__)\n",
126 "\n",
127 "OLLAMA_MODEL = os.getenv(\"OLLAMA_MODEL\", \"mistral:latest\")\n",
128 "ATHENA_DB = os.getenv(\"ATHENA_DATABASE\")\n",
129 "ATHENA_OUTPUT = os.getenv(\"ATHENA_OUTPUT_LOCATION\")\n",
130 "AWS_REGION = os.getenv(\"AWS_REGION\", \"us-east-1\")\n",
131 "\n",
132 "logging.basicConfig(level=logging.INFO)\n",
133 "logger = logging.getLogger(__name__)\n",
134 "\n",
135 "# UI history only (never sent wholesale to the model)\n",
136 "chat_history = []\n",
137 "def init_history():\n",
138 " global chat_history\n",
139 " chat_history = []\n",
140 " logger.info(\"Chat history cleared (UI only).\")\n",
141 "\n",
142 "init_history()\n",
143 "\n",
144 "def run_athena_query(sql: str) -> list:\n",
145 " conn = connect(\n",
146 " s3_staging_dir='s3://ml.dev/athena_query_op/',\n",
147 " region_name='us-east-1'\n",
148 " )\n",
149 " cursor = conn.cursor()\n",
150 " cursor.execute(sql)\n",
151 " cols = [d[0] for d in cursor.description]\n",
152 " rows = cursor.fetchall()\n",
153 " return [dict(zip(cols, row)) for row in rows]\n",
154 "\n",
155 "@app.route(\"/\", methods=[\"GET\"])\n",
156 "def index():\n",
157 " # Render UI with system + history\n",
158 " full = [{\"role\": \"system\", \"content\": SYSTEM_PROMPT}] + chat_history\n",
159 " return render_template(\"index.html\", chat_history=full)\n",
160 "\n",
161 "@app.route(\"/chat\", methods=[\"POST\"])\n",
162 "def chat():\n",
163 " global chat_history\n",
164 " user_text = request.form.get(\"message\", \"\").strip()\n",
165 " logger.info(str(user_text))\n",
166 " if not user_text:\n",
167 " return jsonify(error=\"No message\"), 400\n",
168 "\n",
169 " # Add to UI history\n",
170 " chat_history.append({\"role\": \"user\", \"content\": user_text})\n",
171 "\n",
172 " # ALWAYS send system prompt first, then the *single* user message\n",
173 " payload = [\n",
174 " {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
175 " {\"role\": \"user\", \"content\": user_text}\n",
176 " ]\n",
177 " # logger.info(\"Payload to model: %s\", payload)\n",
178 "\n",
179 " resp = ollama.chat(\n",
180 " model=OLLAMA_MODEL,\n",
181 " messages=payload,\n",
182 " options={\"temperature\": 0}\n",
183 " )\n",
184 " out = resp.message.content\n",
185 " logger.info(\"Model output: %s\", out)\n",
186 "\n",
187 " # Try to parse JSON\n",
188 " try:\n",
189 " result = json.loads(out)\n",
190 " sql = result.get(\"sql_query\", \"\").strip()\n",
191 " except Exception:\n",
192 " # Fallback to raw\n",
193 " chat_history.append({\"role\": \"assistant\", \"content\": out})\n",
194 " full = [{\"role\":\"system\",\"content\":SYSTEM_PROMPT}] + chat_history\n",
195 " return jsonify(chat_history=full)\n",
196 "\n",
197 " # Optionally run it\n",
198 " if sql:\n",
199 " try:\n",
200 " result[\"athena_results\"] = run_athena_query(sql)\n",
201 " except Exception as e:\n",
202 " result[\"error\"] = str(e)\n",
203 "\n",
204 " formatted = json.dumps(result, indent=2)\n",
205 " chat_history.append({\"role\": \"assistant\", \"content\": formatted})\n",
206 "\n",
207 " full = [{\"role\":\"system\",\"content\":SYSTEM_PROMPT}] + chat_history\n",
208 " return jsonify(chat_history=full)\n",
209 "\n",
210 "@app.route(\"/clear\", methods=[\"POST\"])\n",
211 "def clear():\n",
212 " init_history()\n",
213 " full = [{\"role\":\"system\",\"content\":SYSTEM_PROMPT}] + chat_history\n",
214 " return jsonify(chat_history=full)\n",
215 "\n",
216 "if __name__ == \"__main__\":\n",
217 " app.run(debug=False, use_reloader=False)\n"
218 ]
219 },
220 {
221 "cell_type": "code",
222 "execution_count": 4,
223 "id": "12aac119-800f-4306-9a3e-1660f7333a4b",
224 "metadata": {},
225 "outputs": [
226 {
227 "name": "stderr",
228 "output_type": "stream",
229 "text": [
230 "INFO:__main__:Chat history cleared (UI only).\n"
231 ]
232 },
233 {
234 "ename": "TypeError",
235 "evalue": "chat() takes 0 positional arguments but 1 was given",
236 "output_type": "error",
237 "traceback": [
238 "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
239 "\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
240 "\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",
241 "\u001b[0;31mTypeError\u001b[0m: chat() takes 0 positional arguments but 1 was given"
242 ]
243 }
244 ],
245 "source": [
246 "init_history()\n",
247 "chat(\"total crawled products for petco and amazon both in 2025\")"
248 ]
249 },
250 {
251 "cell_type": "code",
252 "execution_count": null,
253 "id": "2782ace6-6234-4990-a920-e15ff9dd3f7f",
254 "metadata": {},
255 "outputs": [],
256 "source": [
257 "session"
258 ]
259 },
260 {
261 "cell_type": "code",
262 "execution_count": null,
263 "id": "a0538594-88bf-4ae9-9510-3ad96e17c8fa",
264 "metadata": {},
265 "outputs": [],
266 "source": [
267 "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'))])"
268 ]
269 },
270 {
271 "cell_type": "code",
272 "execution_count": 5,
273 "id": "1c84683f-67d7-4316-88ae-a2e1680c85db",
274 "metadata": {},
275 "outputs": [
276 {
277 "data": {
278 "text/plain": [
279 "[]"
280 ]
281 },
282 "execution_count": 5,
283 "metadata": {},
284 "output_type": "execute_result"
285 }
286 ],
287 "source": [
288 "chat_history"
289 ]
290 },
291 {
292 "cell_type": "code",
293 "execution_count": 7,
294 "id": "7d166f3f-7656-45ad-b196-9c340002637a",
295 "metadata": {},
296 "outputs": [
297 {
298 "data": {
299 "application/vnd.jupyter.widget-view+json": {
300 "model_id": "f9dc06df4fc44ff49ac3552b1ddf0db0",
301 "version_major": 2,
302 "version_minor": 0
303 },
304 "text/plain": [
305 "Loading pipeline components...: 0%| | 0/7 [00:00<?, ?it/s]"
306 ]
307 },
308 "metadata": {},
309 "output_type": "display_data"
310 },
311 {
312 "name": "stderr",
313 "output_type": "stream",
314 "text": [
315 "You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers\n"
316 ]
317 },
318 {
319 "data": {
320 "application/vnd.jupyter.widget-view+json": {
321 "model_id": "ee8073a0e92a486db7880c6f8484b4ac",
322 "version_major": 2,
323 "version_minor": 0
324 },
325 "text/plain": [
326 "Loading checkpoint shards: 0%| | 0/3 [00:00<?, ?it/s]"
327 ]
328 },
329 "metadata": {},
330 "output_type": "display_data"
331 },
332 {
333 "data": {
334 "application/vnd.jupyter.widget-view+json": {
335 "model_id": "a0ce727d4c174defa6dbea5a9f3f93df",
336 "version_major": 2,
337 "version_minor": 0
338 },
339 "text/plain": [
340 "Loading checkpoint shards: 0%| | 0/2 [00:00<?, ?it/s]"
341 ]
342 },
343 "metadata": {},
344 "output_type": "display_data"
345 },
346 {
347 "ename": "OutOfMemoryError",
348 "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)",
349 "output_type": "error",
350 "traceback": [
351 "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
352 "\u001b[0;31mOutOfMemoryError\u001b[0m Traceback (most recent call last)",
353 "\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",
354 "\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",
355 "\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",
356 "\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",
357 "\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",
358 "\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",
359 "\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",
360 "\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",
361 "\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",
362 "\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",
363 "\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",
364 "\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",
365 "\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",
366 "\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)"
367 ]
368 }
369 ],
370 "source": [
371 "from diffusers import DiffusionPipeline\n",
372 "import torch\n",
373 "\n",
374 "# quick test\n",
375 "pipe = DiffusionPipeline.from_pretrained(\"black-forest-labs/FLUX.1-dev\")\n",
376 "pipe.to(\"cuda\")\n",
377 "img = pipe(\"Astronaut in a jungle, cold color palette, detailed, 8k\").images[0]\n",
378 "img.save(\"test.png\")\n",
379 "print(\"✅ Generated test.png\")\n"
380 ]
381 },
382 {
383 "cell_type": "code",
384 "execution_count": null,
385 "id": "46d533a1-c05a-43d9-95e5-3a24bc196a50",
386 "metadata": {},
387 "outputs": [],
388 "source": [
389 "# In your terminal (or in a Jupyter cell prefixing with !)\n",
390 "!pip install diffusers transformers accelerate safetensors\n"
391 ]
392 },
393 {
394 "cell_type": "code",
395 "execution_count": null,
396 "id": "1be04254-d787-458e-963a-e379d179cafa",
397 "metadata": {},
398 "outputs": [],
399 "source": []
400 },
401 {
402 "cell_type": "code",
403 "execution_count": null,
404 "id": "ffdcc284-a56c-45ab-b8bc-b88ab0609fa0",
405 "metadata": {},
406 "outputs": [],
407 "source": []
408 }
409 ],
410 "metadata": {
411 "kernelspec": {
412 "display_name": "Python 3",
413 "language": "python",
414 "name": "python3"
415 },
416 "language_info": {
417 "codemirror_mode": {
418 "name": "ipython",
419 "version": 3
420 },
421 "file_extension": ".py",
422 "mimetype": "text/x-python",
423 "name": "python",
424 "nbconvert_exporter": "python",
425 "pygments_lexer": "ipython3",
426 "version": "3.8.10"
427 }
428 },
429 "nbformat": 4,
430 "nbformat_minor": 5
431 }
432