Update

colab.ipynb

@@ -13,12 +13,11 @@
     "import ipywidgets as widgets\n",
     "from huggingface_hub import InferenceClient\n",
     "from IPython.display import display, clear_output\n",
-    "from img_gen_logic import generate_image\n",
+    "from img_gen_logic_colab import generate_image, save_image\n",
     "from config_colab import models, prompts, api_token\n",
     "from PIL import Image\n",
     "from google.colab import userdata\n",
     "from datetime import datetime\n",
-    "from img_gen_logic_colab import generate_image, save_image\n",
     "\n",
     "# Initialize the InferenceClient with the default model\n",
     "client = InferenceClient(models[0][\"name\"], token=api_token)\n",
@@ -110,6 +109,7 @@
     "# Output area to display the image\n",
     "output = widgets.Output()\n",
     "\n",
+    "# Function to handle button click event\n",
     "def on_generate_button_clicked(b):\n",
     "    with output:\n",
     "        clear_output(wait=True)  # Clear previous output\n",
@@ -142,30 +142,30 @@
     "        print(\"=== Debug: Calling generate_image ===\")\n",
     "        image, message = generate_image(\n",
     "            selected_prompt, selected_team, selected_model, height, width,\n",
-    "            num_inference_steps, guidance_scale, seed, custom_prompt, api_token\n",
+    "            num_inference_steps, guidance_scale, seed, custom_prompt, api_token,\n",
+    "            randomize_seed=randomize_seed_checkbox.value\n",
     "        )\n",
     "\n",
     "        # Debug: Check the output of generate_image\n",
     "        print(\"=== Debug: generate_image Output ===\")\n",
-    "        print(f\"Image: {image}\")\n",
     "        print(f\"Message: {message}\")\n",
     "        print(\"====================================\")\n",
     "\n",
-    "        if isinstance(image, str):\n",
-    "            print(\"=== Debug: Error ===\")\n",
-    "            print(image)\n",
-    "        else:\n",
+    "        if image is not None:\n",
     "            # Debug: Indicate that the image is being displayed and saved\n",
     "            print(\"=== Debug: Image Generation ===\")\n",
     "            print(\"Image generated successfully!\")\n",
     "            print(\"Displaying image...\")\n",
     "\n",
     "            # Display the image in the notebook\n",
-    "            if image is not None:\n",
-    "                display(image)\n",
-    "            else:\n",
-    "                print(\"=== Debug: Error ===\")\n",
-    "                print(\"No image was returned by generate_image.\")\n",
+    "            display(image)\n",
+    "\n",
+    "            # Save the image with a timestamped filename\n",
+    "            output_filename = save_image(image, model_dropdown.label, prompt_dropdown.label, selected_team)\n",
+    "            print(f\"Image saved as {output_filename}\")\n",
+    "        else:\n",
+    "            print(\"=== Debug: Error ===\")\n",
+    "            print(message)\n",
     "\n",
     "# Attach the button click event handler\n",
     "generate_button.on_click(on_generate_button_clicked)\n",

img_gen_logic_colab.py

@@ -4,7 +4,6 @@ from PIL import Image
 import random
 from datetime import datetime
 
-# Function to generate images based on the selected prompt, team, and model
 def generate_image(prompt, team, model_name, height, width, num_inference_steps, guidance_scale, seed, custom_prompt, api_token, randomize_seed=True):
     """
     Generate an image using the Hugging Face Inference API.
@@ -23,7 +22,7 @@ def generate_image(prompt, team, model_name, height, width, num_inference_steps,
         randomize_seed (bool): Whether to randomize the seed.
 
     Returns:
-        PIL.Image.Image or str: The generated image or an error message.
+        tuple: A tuple containing the generated image (PIL.Image.Image) and a message (str).
     """
     # Determine the enemy color
     enemy_color = "blue" if team.lower() == "red" else "red"
@@ -37,7 +36,7 @@ def generate_image(prompt, team, model_name, height, width, num_inference_steps,
     elif team.lower() == "blue":
         prompt += " The winning army is dressed in blue armor and banners."
     else:
-        return "Invalid team selection. Please choose 'Red' or 'Blue'."
+        return None, "Invalid team selection. Please choose 'Red' or 'Blue'."
 
     # Append the custom prompt if provided
     if custom_prompt.strip():
@@ -65,11 +64,10 @@ def generate_image(prompt, team, model_name, height, width, num_inference_steps,
             height=height,  # Height
             seed=seed  # Random seed
         )
-        return image
+        return image, "Image generated successfully."
     except Exception as e:
-        return f"An error occurred: {e}"
+        return None, f"An error occurred: {e}"
 
-# Function to save the generated image with a timestamped filename
 def save_image(image, model_label, prompt_label, team):
     """
     Save the generated image with a timestamped filename.