radames commited on
Commit
85c974d
β€’
1 Parent(s): c211fda

throw on nsfw palette

Browse files
Files changed (1) hide show
  1. frontend/src/routes/+page.svelte +19 -6
frontend/src/routes/+page.svelte CHANGED
@@ -97,12 +97,16 @@
97
  $loadingState = data.success ? 'Generating' : 'Error';
98
  break;
99
  case 'process_completed':
100
- $loadingState = data.success ? 'Complete' : 'Error';
101
- const images = await extractColorsImages(data.output.data[0], _prompt);
102
- savePaletteDB({
103
- prompt: _prompt,
104
- images
105
- });
 
 
 
 
106
  websocket.close();
107
  $isLoading = false;
108
  return;
@@ -118,9 +122,18 @@
118
  };
119
  }
120
  async function extractColorsImages(images: string[], _prompt: string): Promise<ColorsImage[]> {
 
121
  const colorImages = [];
122
  for (const base64img of images) {
123
  const { colors, imgBlob } = await extractPalette(base64img);
 
 
 
 
 
 
 
 
124
  const url = await uploadImage(imgBlob, _prompt);
125
  const colorsImage: ColorsImage = {
126
  colors,
 
97
  $loadingState = data.success ? 'Generating' : 'Error';
98
  break;
99
  case 'process_completed':
100
+ try {
101
+ const images = await extractColorsImages(data.output.data[0], _prompt);
102
+ savePaletteDB({
103
+ prompt: _prompt,
104
+ images
105
+ });
106
+ $loadingState = data.success ? 'Complete' : 'Error';
107
+ } catch (e) {
108
+ $loadingState = e.message;
109
+ }
110
  websocket.close();
111
  $isLoading = false;
112
  return;
 
122
  };
123
  }
124
  async function extractColorsImages(images: string[], _prompt: string): Promise<ColorsImage[]> {
125
+ const nsfwColors = ['#040404', '#b7b7b7', '#565656', '#747474', '#6c6c6c'];
126
  const colorImages = [];
127
  for (const base64img of images) {
128
  const { colors, imgBlob } = await extractPalette(base64img);
129
+ if (
130
+ !colors
131
+ .map((color) => color.formatHex().toUpperCase())
132
+ .every((c) => nsfwColors.includes(c.toUpperCase()))
133
+ ) {
134
+ console.error('Possible NSFW image');
135
+ throw new Error('Possible NSFW image');
136
+ }
137
  const url = await uploadImage(imgBlob, _prompt);
138
  const colorsImage: ColorsImage = {
139
  colors,