Update

colab.ipynb

@@ -109,7 +109,6 @@
     "# 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",
@@ -140,7 +139,7 @@
     "\n",
     "        # Generate the image\n",
     "        print(\"=== Debug: Calling generate_image ===\")\n",
-    "        image, message = generate_image(\n",
+    "        image = generate_image(\n",
     "            selected_prompt, selected_team, selected_model, height, width,\n",
     "            num_inference_steps, guidance_scale, seed, custom_prompt, api_token,\n",
     "            randomize_seed=randomize_seed_checkbox.value\n",
@@ -148,10 +147,13 @@
     "\n",
     "        # Debug: Check the output of generate_image\n",
     "        print(\"=== Debug: generate_image Output ===\")\n",
-    "        print(f\"Message: {message}\")\n",
+    "        print(f\"Image: {image}\")\n",
     "        print(\"====================================\")\n",
     "\n",
-    "        if image is not None:\n",
+    "        if isinstance(image, str):\n",
+    "            print(\"=== Debug: Error ===\")\n",
+    "            print(image)\n",
+    "        else:\n",
     "            # Debug: Indicate that the image is being displayed and saved\n",
     "            print(\"=== Debug: Image Generation ===\")\n",
     "            print(\"Image generated successfully!\")\n",
@@ -163,10 +165,7 @@
     "            # 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",
+    "            \n",
     "# Attach the button click event handler\n",
     "generate_button.on_click(on_generate_button_clicked)\n",
     "\n",

img_gen_logic_colab.py

@@ -22,7 +22,7 @@ def generate_image(prompt, team, model_name, height, width, num_inference_steps,
         randomize_seed (bool): Whether to randomize the seed.
 
     Returns:
-        tuple: A tuple containing the generated image (PIL.Image.Image) and a message (str).
+        PIL.Image.Image or str: The generated image or an error message.
     """
     # Determine the enemy color
     enemy_color = "blue" if team.lower() == "red" else "red"
@@ -36,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 None, "Invalid team selection. Please choose 'Red' or 'Blue'."
+        return "Invalid team selection. Please choose 'Red' or 'Blue'."
 
     # Append the custom prompt if provided
     if custom_prompt.strip():
@@ -64,9 +64,9 @@ def generate_image(prompt, team, model_name, height, width, num_inference_steps,
             height=height,  # Height
             seed=seed  # Random seed
         )
-        return image, "Image generated successfully."
+        return image
     except Exception as e:
-        return None, f"An error occurred: {e}"
+        return f"An error occurred: {e}"
 
 def save_image(image, model_label, prompt_label, team):
     """