colab.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'Python 3.12.0' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: '/usr/local/bin/python3.12 -m pip install ipykernel -U --user --force-reinstall'"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# Set PYTHONPATH to the project root \n",
    "# Solves all problems w subfolders - option1\n",
    "os.environ[\"PYTHONPATH\"] = os.path.abspath(os.path.join(\"..\"))\n",
    "\n",
    "import random\n",
    "from huggingface_hub import InferenceClient\n",
    "from PIL import Image\n",
    "from google.colab import userdata\n",
    "from IPython.display import display, clear_output\n",
    "import ipywidgets as widgets\n",
    "from datetime import datetime\n",
    "from config.config_colab import api_token\n",
    "from config.models import models\n",
    "from config.prompts import generate_prompt\n",
    "\n",
    "# Initialize the InferenceClient with the default model\n",
    "client = InferenceClient(models[0][\"name\"], token=api_token)\n",
    "\n",
    "\n",
    "# Input for left castle HP\n",
    "left_hp_input = widgets.IntSlider(\n",
    "    value=100,\n",
    "    min=0,\n",
    "    max=100,\n",
    "    step=1,\n",
    "    description=\"Left Castle HP:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for right castle HP\n",
    "right_hp_input = widgets.IntSlider(\n",
    "    value=100,\n",
    "    min=0,\n",
    "    max=100,\n",
    "    step=1,\n",
    "    description=\"Right Castle HP:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for height\n",
    "height_input = widgets.IntText(\n",
    "    value=512,\n",
    "    description=\"Height:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for width\n",
    "width_input = widgets.IntText(\n",
    "    value=1024,\n",
    "    description=\"Width:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for number of inference steps\n",
    "num_inference_steps_input = widgets.IntSlider(\n",
    "    value=20,\n",
    "    min=10,\n",
    "    max=100,\n",
    "    step=1,\n",
    "    description=\"Inference Steps:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for guidance scale (default set to 2)\n",
    "guidance_scale_input = widgets.FloatSlider(\n",
    "    value=2.0,  # Default set to 2\n",
    "    min=1.0,\n",
    "    max=20.0,\n",
    "    step=0.5,\n",
    "    description=\"Guidance Scale:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for seed\n",
    "seed_input = widgets.IntText(\n",
    "    value=random.randint(0, 1000000),\n",
    "    description=\"Seed:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Checkbox to randomize seed\n",
    "randomize_seed_checkbox = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description=\"Randomize Seed\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Button to generate image\n",
    "generate_button = widgets.Button(\n",
    "    description=\"Generate Image\",\n",
    "    button_style=\"success\"\n",
    ")\n",
    "\n",
    "\n",
    "# Output area to display the image\n",
    "output = widgets.Output()\n",
    "\n",
    "\n",
    "# Function to generate images based on the HP values\n",
    "def generate_image(left_hp, right_hp, height, width, num_inference_steps, guidance_scale, seed):\n",
    "    # Generate the prompt\n",
    "    prompt = generate_prompt(left_hp, right_hp)\n",
    "\n",
    "    try:\n",
    "        # Randomize the seed if the checkbox is checked\n",
    "        if randomize_seed_checkbox.value:\n",
    "            seed = random.randint(0, 1000000)\n",
    "            seed_input.value = seed  # Update the seed input box\n",
    "\n",
    "        print(f\"Using seed: {seed}\")\n",
    "\n",
    "        # Debug: Indicate that the image is being generated\n",
    "        print(\"Generating image... Please wait.\")\n",
    "\n",
    "        # Initialize the InferenceClient with the selected model\n",
    "        client = InferenceClient(models[0][\"name\"], token=api_token)\n",
    "\n",
    "        # Generate the image using the Inference API with parameters\n",
    "        image = client.text_to_image(\n",
    "            prompt,\n",
    "            guidance_scale=guidance_scale,  # Guidance scale\n",
    "            num_inference_steps=num_inference_steps,  # Number of inference steps\n",
    "            width=width,  # Width\n",
    "            height=height,  # Height\n",
    "            seed=seed  # Random seed\n",
    "        )\n",
    "        return image\n",
    "    except Exception as e:\n",
    "        return f\"An error occurred: {e}\"\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",
    "        left_hp = left_hp_input.value\n",
    "        right_hp = right_hp_input.value\n",
    "        height = height_input.value\n",
    "        width = width_input.value\n",
    "        num_inference_steps = num_inference_steps_input.value\n",
    "        guidance_scale = guidance_scale_input.value\n",
    "        seed = seed_input.value\n",
    "\n",
    "        # Debug: Show selected parameters\n",
    "        print(f\"Left Castle HP: {left_hp}\")\n",
    "        print(f\"Right Castle HP: {right_hp}\")\n",
    "        print(f\"Height: {height}\")\n",
    "        print(f\"Width: {width}\")\n",
    "        print(f\"Inference Steps: {num_inference_steps}\")\n",
    "        print(f\"Guidance Scale: {guidance_scale}\")\n",
    "        print(f\"Seed: {seed}\")\n",
    "\n",
    "        # Generate the image\n",
    "        image = generate_image(left_hp, right_hp, height, width, num_inference_steps, guidance_scale, seed)\n",
    "\n",
    "        if isinstance(image, str):\n",
    "            print(image)\n",
    "        else:\n",
    "            # Debug: Indicate that the image is being displayed and saved\n",
    "            print(\"Image generated successfully!\")\n",
    "            print(\"Displaying image...\")\n",
    "\n",
    "            # Display the image in the notebook\n",
    "            display(image)\n",
    "\n",
    "            # Save the image with a timestamped filename\n",
    "            timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "            output_filename = f\"{timestamp}_left_{left_hp}_right_{right_hp}.png\"\n",
    "            print(f\"Saving image as {output_filename}...\")\n",
    "            image.save(output_filename)\n",
    "            print(f\"Image saved as {output_filename}\")\n",
    "\n",
    "\n",
    "# Attach the button click event handler\n",
    "generate_button.on_click(on_generate_button_clicked)\n",
    "\n",
    "# Display the widgets\n",
    "#display(left_hp_input, right_hp_input, height_input, width_input, num_inference_steps_input, guidance_scale_input, seed_input, randomize_seed_checkbox, generate_button, output)\n",
    "\n",
    "display(left_hp_input, right_hp_input, generate_button, output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}