From a937f6504a6a4e692acdd7a33ee83f533c52783e Mon Sep 17 00:00:00 2001
From: Jun Siang Cheah <git@jscheah.me>
Date: Mon, 20 May 2024 08:56:18 +0100
Subject: [PATCH] feat: add image gen with automatic1111 to integration test

---
 .github/workflows/integration-test.yml |   6 ++++-
 cypress/e2e/chat.cy.ts                 |  23 ++++++++++++++++++
 docker-compose.a1111-test.yaml         |  31 +++++++++++++++++++++++++
 src/lib/components/common/Image.svelte |   2 +-
 test/test_files/image_gen/sd-empty.pt  | Bin 0 -> 431 bytes
 5 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 docker-compose.a1111-test.yaml
 create mode 100644 test/test_files/image_gen/sd-empty.pt

diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml
index 32346d3b9..96ba50289 100644
--- a/.github/workflows/integration-test.yml
+++ b/.github/workflows/integration-test.yml
@@ -20,7 +20,11 @@ jobs:
 
       - name: Build and run Compose Stack
         run: |
-          docker compose --file docker-compose.yaml --file docker-compose.api.yaml up --detach --build
+          docker compose \
+            --file docker-compose.yaml \
+            --file docker-compose.api.yaml \
+            --file docker-compose.a1111-test.yaml \
+            up --detach --build
           
       - name: Wait for Ollama to be up
         timeout-minutes: 5
diff --git a/cypress/e2e/chat.cy.ts b/cypress/e2e/chat.cy.ts
index ced998104..ddb33d6c0 100644
--- a/cypress/e2e/chat.cy.ts
+++ b/cypress/e2e/chat.cy.ts
@@ -74,5 +74,28 @@ describe('Settings', () => {
 				expect(spy).to.be.callCount(2);
 			});
 		});
+
+		it('user can generate image', () => {
+			// Click on the model selector
+			cy.get('button[aria-label="Select a model"]').click();
+			// Select the first model
+			cy.get('button[aria-label="model-item"]').first().click();
+			// Type a message
+			cy.get('#chat-textarea').type('Hi, what can you do? A single sentence only please.', {
+				force: true
+			});
+			// Send the message
+			cy.get('button[type="submit"]').click();
+			// User's message should be visible
+			cy.get('.chat-user').should('exist');
+			// Wait for the response
+			cy.get('.chat-assistant', { timeout: 120_000 }) // .chat-assistant is created after the first token is received
+				.find('div[aria-label="Generation Info"]', { timeout: 120_000 }) // Generation Info is created after the stop token is received
+				.should('exist');
+			// Click on the generate image button
+			cy.get('[aria-label="Generate Image"]').click();
+			// Wait for image to be visible
+			cy.get('img[data-cy="image"]', { timeout: 60_000 }).should('be.visible');
+		});
 	});
 });
diff --git a/docker-compose.a1111-test.yaml b/docker-compose.a1111-test.yaml
new file mode 100644
index 000000000..e6ab12c07
--- /dev/null
+++ b/docker-compose.a1111-test.yaml
@@ -0,0 +1,31 @@
+# This is an overlay that spins up stable-diffusion-webui for integration testing
+# This is not designed to be used in production
+services:
+  stable-diffusion-webui:
+    # Not built for ARM64
+    platform: linux/amd64
+    image: ghcr.io/neggles/sd-webui-docker:latest
+    restart: unless-stopped
+    environment:
+      CLI_ARGS: "--api --use-cpu all --precision full --no-half --skip-torch-cuda-test --ckpt /empty.pt --do-not-download-clip --disable-nan-check --disable-opt-split-attention"
+      PYTHONUNBUFFERED: "1"
+      TERM: "vt100"
+      SD_WEBUI_VARIANT: "default"
+    # Hack to get container working on Apple Silicon
+    # Rosetta creates a conflict ${HOME}/.cache folder
+    entrypoint: /bin/bash
+    command:
+      - -c
+      - |
+        export HOME=/root-home
+        rm -rf $${HOME}/.cache
+        /docker/entrypoint.sh python -u webui.py --listen --port $${WEBUI_PORT} --skip-version-check $${CLI_ARGS}
+    volumes:
+      - ./test/test_files/image_gen/sd-empty.pt:/empty.pt
+
+  open-webui:
+    environment:
+      ENABLE_IMAGE_GENERATION: "true"
+      AUTOMATIC1111_BASE_URL: http://stable-diffusion-webui:7860
+      IMAGE_SIZE: "64x64"
+      IMAGE_STEPS: "3"
diff --git a/src/lib/components/common/Image.svelte b/src/lib/components/common/Image.svelte
index e69f0e29b..dc65df87d 100644
--- a/src/lib/components/common/Image.svelte
+++ b/src/lib/components/common/Image.svelte
@@ -19,5 +19,5 @@
 		showImagePreview = true;
 	}}
 >
-	<img src={_src} {alt} class=" max-h-96 rounded-lg" draggable="false" />
+	<img src={_src} {alt} class=" max-h-96 rounded-lg" draggable="false" data-cy="image" />
 </button>
diff --git a/test/test_files/image_gen/sd-empty.pt b/test/test_files/image_gen/sd-empty.pt
new file mode 100644
index 0000000000000000000000000000000000000000..c6ac59eb01fcb778290a85f12bdb7867de3dfdd1
GIT binary patch
literal 431
zcmWIWW@cev;NW1u00Im`42ea_8JT6N`YDMeiFyUuIc`pT3{fbcfvL8TK`+3Yo#WF)
zrlV{?Q$RQXr>Xo5ws2F+Qj3Z+^Yh%CEYS=_u>n8Fm<we0#d+zWOh9`<IKZ0`M8W-p
z>@ES21PVa+A-Zm4llf6}h5>mn-B6zdc(bwTKo!X`>%x_T-2>#o=xV6UB`6Kl#|~op
WGC~AERDd@tC?tV;m>59nA!-3{+(-BT

literal 0
HcmV?d00001