Update

colab.ipynb

@@ -6,16 +6,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# colab.ipynb\n",
     "# Import necessary libraries\n",
     "import random\n",
     "import ipywidgets as widgets\n",
     "from IPython.display import display, clear_output\n",
     "from img_gen_logic import generate_image\n",
     "from config import models, prompts\n",
-    "from google.colab import userdata\n",
-    "\n",
-    "# Retrieve the Hugging Face token from Colab secrets\n",
-    "api_token = userdata.get(\"HF_CTB_TOKEN\")\n",
     "\n",
     "# Dropdown menu for model selection\n",
     "model_dropdown = widgets.Dropdown(\n",

config.py

@@ -1,15 +1,22 @@
-# src/config.py
+# config.py
 import os
+from google.colab import userdata  # Only for Colab
 
-# Retrieve the Hugging Face token from environment variables
-api_token = os.getenv("HF_CTB_TOKEN")
+# Retrieve the Hugging Face token
+try:
+    # Try to get the token from Colab secrets
+    api_token = userdata.get("HF_CTB_TOKEN")
+except ImportError:
+    # Fall back to environment variable (for local execution)
+    api_token = os.getenv("HF_CTB_TOKEN")
 
 # Debugging: Check if the Hugging Face token is available
 if not api_token:
-    print("ERROR: Hugging Face token (HF_CTB_TOKEN) is missing. Please set it as an environment variable.")
+    print("ERROR: Hugging Face token (HF_CTB_TOKEN) is missing. Please set it as an environment variable or in Colab secrets.")
 else:
     print("Hugging Face token loaded successfully.")
 
+
 # List of models with aliases
 models = [
     {"alias": "FLUX.1-dev", "name": "black-forest-labs/FLUX.1-dev"},