diff --git a/README.md b/README.md
index de6cc2e..30564e2 100644
--- a/README.md
+++ b/README.md
@@ -113,6 +113,8 @@ Then browse to
`http://localhost:7860/?__theme=dark`
+
+
Optionally, you can use the following command-line flags:
| Flag | Description |
@@ -121,7 +123,7 @@ Optionally, you can use the following command-line flags:
| `--model MODEL` | Name of the model to load by default. |
| `--notebook` | Launch the web UI in notebook mode, where the output is written to the same text box as the input. |
| `--chat` | Launch the web UI in chat mode.|
-| `--cai-chat` | Launch the web UI in chat mode with a style similar to Character.AI's. If the file profile.png or profile.jpg exists in the same folder as server.py, this image will be used as the bot's profile picture. |
+| `--cai-chat` | Launch the web UI in chat mode with a style similar to Character.AI's. If the file img-bot.png or img-bot.jpg exists in the same folder as server.py, this image will be used as the bot's profile picture. Similarly, img-you.png or img-you.jpg will be used as your profile picture. |
| `--cpu` | Use the CPU to generate text.|
| `--load-in-8bit` | Load the model with 8-bit precision.|
| `--auto-devices` | Automatically split the model across the available GPU(s) and CPU.|
diff --git a/modules/html_generator.py b/modules/html_generator.py
index e5b6f82..11da068 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -196,7 +196,7 @@ def generate_chat_html(history, name1, name2, character):
border-radius: 50%;
}
- .circle-bot img {
+ .circle-bot img, .circle-you img {
border-radius: 50%;
width: 100%;
height: 100%;
@@ -225,15 +225,21 @@ def generate_chat_html(history, name1, name2, character):
f"characters/{character}.png",
f"characters/{character}.jpg",
f"characters/{character}.jpeg",
- "profile.png",
- "profile.jpg",
- "profile.jpeg",
+ "img_bot.png",
+ "img_bot.jpg",
+ "img_bot.jpeg"
]:
if Path(i).exists():
img = f''
break
+ img_you = ''
+ for i in ["img_me.png", "img_me.jpg", "img_me.jpeg"]:
+ if Path(i).exists():
+ img_you = f'
'
+ break
+
for i,_row in enumerate(history[::-1]):
row = _row.copy()
row[0] = re.sub(r"[\\]*\*", r"*", row[0])
@@ -262,6 +268,7 @@ def generate_chat_html(history, name1, name2, character):
output += f"""