File size: 5,982 Bytes
e1eaf41
 
 
 
 
 
 
 
96410e5
e1eaf41
465c4f7
e1eaf41
6dd0a15
e1eaf41
 
 
 
0f4c106
 
 
e1eaf41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# colab.ipynb\n",
    "# Import necessary libraries\n",
    "import random\n",
    "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 config import models, prompts\n",
    "\n",
    "# Initialize the InferenceClient with the default model\n",
    "client = InferenceClient(models[0][\"name\"], token=api_token)\n",
    "\n",
    "# Dropdown menu for model selection\n",
    "model_dropdown = widgets.Dropdown(\n",
    "    options=[(model[\"alias\"], model[\"name\"]) for model in models],\n",
    "    description=\"Select Model:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Dropdown menu for prompt selection\n",
    "prompt_dropdown = widgets.Dropdown(\n",
    "    options=[(prompt[\"alias\"], prompt[\"text\"]) for prompt in prompts],\n",
    "    description=\"Select Prompt:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Dropdown menu for team selection\n",
    "team_dropdown = widgets.Dropdown(\n",
    "    options=[\"Red\", \"Blue\"],\n",
    "    description=\"Select Team:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for height\n",
    "height_input = widgets.IntText(\n",
    "    value=360,\n",
    "    description=\"Height:\",\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "# Input for width\n",
    "width_input = widgets.IntText(\n",
    "    value=640,\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\n",
    "guidance_scale_input = widgets.FloatSlider(\n",
    "    value=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",
    "# Text box for custom prompt\n",
    "custom_prompt_input = widgets.Textarea(\n",
    "    value=\"\",\n",
    "    placeholder=\"Enter your custom prompt (up to 200 characters)...\",\n",
    "    description=\"Custom Prompt:\",\n",
    "    style={\"description_width\": \"initial\"},\n",
    "    layout=widgets.Layout(width=\"500px\", height=\"80px\")\n",
    ")\n",
    "\n",
    "# Button to generate image\n",
    "generate_button = widgets.Button(\n",
    "    description=\"Generate Image\",\n",
    "    button_style=\"success\"\n",
    ")\n",
    "\n",
    "# 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",
    "        selected_prompt = prompt_dropdown.value\n",
    "        selected_team = team_dropdown.value\n",
    "        selected_model = model_dropdown.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",
    "        custom_prompt = custom_prompt_input.value\n",
    "\n",
    "        # Debug: Show selected parameters\n",
    "        print(f\"Selected Model: {model_dropdown.label}\")\n",
    "        print(f\"Selected Prompt: {prompt_dropdown.label}\")\n",
    "        print(f\"Selected Team: {selected_team}\")\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",
    "        print(f\"Custom Prompt: {custom_prompt}\")\n",
    "\n",
    "        # Generate the image\n",
    "        image, message = generate_image(selected_prompt, selected_team, selected_model, height, width, num_inference_steps, guidance_scale, seed, custom_prompt)\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",
    "# Attach the button click event handler\n",
    "generate_button.on_click(on_generate_button_clicked)\n",
    "\n",
    "# Display the widgets\n",
    "display(prompt_dropdown, team_dropdown, model_dropdown, custom_prompt_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
}