Merge branch 'main' of https://github.com/Dungeon-MASSters/MASSter

Dungeon-MASSters · Oct 14, 2023 · 70febcb · 70febcb
2 parents 4b7d783 + 4189789
commit 70febcb
Show file tree

Hide file tree

Showing 22 changed files with 921 additions and 50 deletions.
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -0,0 +1,46 @@
+version: '3'
+
+volumes:
+  pb_volume:
+
+services:
+  pocketbase:
+    build:
+      context: ./pocketbase
+
+    ports:
+      - 8090:8090
+
+    command: ["/pocketbase/pocketbase", "serve", "--http", "0.0.0.0:8090"]
+
+    volumes:
+      - pb_volume:/pocketbase
+
+  frontend-dev:
+    build:
+      context: ./frontend/MASSter-frontend
+
+    ports:
+      - 8080:8080
+
+    command: ["npm", "run", "preview", "--", "--host", "0.0.0.0", "--port", "8080"]
+
+  preprocessor:
+    build:
+      context: model/preprocessor/docker
+
+    environment:
+      - PB_LINK=http://pocketbase:8090/
+      - [email protected]
+      - PB_PWD=1234567890
+      - COL_NAME=text_generation_mvp
+
+  generator:
+    build:
+      context: model/generator/docker
+
+    environment:
+      - PB_LINK=http://pocketbase:8090/
+      - [email protected]
+      - PB_PWD=1234567890
+      - COL_NAME=text_generation_mvp
diff --git a/frontend/MASSter-frontend/.dockerignore b/frontend/MASSter-frontend/.dockerignore
@@ -0,0 +1,2 @@
+node_modules/
+dist/
diff --git a/frontend/MASSter-frontend/.env.production b/frontend/MASSter-frontend/.env.production
@@ -0,0 +1 @@
+VITE_PB_HOST="http://pocketbase:8090"
diff --git a/frontend/MASSter-frontend/Dockerfile b/frontend/MASSter-frontend/Dockerfile
@@ -0,0 +1,12 @@
+FROM node:18
+
+COPY .env.production .
+COPY package.json .
+COPY package-lock.json .
+
+RUN npm install
+
+COPY . .
+
+RUN npm run build
+
diff --git a/frontend/MASSter-frontend/src/components/avatar-prompt.tsx b/frontend/MASSter-frontend/src/components/avatar-prompt.tsx
@@ -171,17 +171,13 @@ export function AvatarAdvancedPromptForm({
         negativePrompt: z.string().max(500).optional(),
         style: z.string().default("Нет"),
         inputImage: z
-            .custom<FileList>((v) => v instanceof FileList, {
+            .custom<File>((v) => v instanceof File, {
                 message: "Нужно загрузить картинки для референса"
             })
-            .refine((files) => {
-                for (let file of files) {
-                    if (!ACCEPTED_IMAGE_TYPES.includes(file.type)) {
-                        return false;
-                    }
-                }
-                return true;
-            }, "Пока мы поддерживаем только .png и .jpg изображения")
+            .refine(
+                (v) => ACCEPTED_IMAGE_TYPES.includes(v.type),
+                "Пока мы поддерживаем только .png и .jpg изображения"
+            )
             .optional(),
         // video: z
         //     .custom<File>((v) => v instanceof File, {
@@ -213,9 +209,7 @@ export function AvatarAdvancedPromptForm({
         formData.append("num_images", `${data.numImages}`);
 
         if (data.inputImage) {
-            for (let file of data.inputImage) {
-                formData.append("input_image", file);
-            }
+            formData.append("input_image", data.inputImage);
         }
         // if (data.video) {
         //     formData.append("video", data.video);
@@ -347,7 +341,6 @@ export function AvatarAdvancedPromptForm({
                                             <FormControl>
                                                 <Input
                                                     type="file"
-                                                    multiple={true}
                                                     onBlur={field.onBlur}
                                                     onChange={(e) =>
                                                         field.onChange(

diff --git a/frontend/MASSter-frontend/src/components/banner-prompt.tsx b/frontend/MASSter-frontend/src/components/banner-prompt.tsx
@@ -171,17 +171,13 @@ export function BannerAdvancedPromptForm({
         negativePrompt: z.string().max(500).optional(),
         style: z.string().default("Нет"),
         inputImage: z
-            .custom<FileList>((v) => v instanceof FileList, {
+            .custom<File>((v) => v instanceof File, {
                 message: "Нужно загрузить картинки для референса"
             })
-            .refine((files) => {
-                for (let file of files) {
-                    if (!ACCEPTED_IMAGE_TYPES.includes(file.type)) {
-                        return false;
-                    }
-                }
-                return true;
-            }, "Пока мы поддерживаем только .png и .jpg изображения")
+            .refine(
+                (v) => ACCEPTED_IMAGE_TYPES.includes(v.type),
+                "Пока мы поддерживаем только .png и .jpg изображения"
+            )
             .optional(),
         // video: z
         //     .custom<File>((v) => v instanceof File, {
@@ -213,9 +209,7 @@ export function BannerAdvancedPromptForm({
         formData.append("num_images", `${data.numImages}`);
 
         if (data.inputImage) {
-            for (let file of data.inputImage) {
-                formData.append("input_image", file);
-            }
+            formData.append("input_image", data.inputImage);
         }
         // if (data.video) {
         //     formData.append("video", data.video);
@@ -347,7 +341,6 @@ export function BannerAdvancedPromptForm({
                                             <FormControl>
                                                 <Input
                                                     type="file"
-                                                    multiple={true}
                                                     onBlur={field.onBlur}
                                                     onChange={(e) =>
                                                         field.onChange(

diff --git a/frontend/MASSter-frontend/src/components/video-cover-prompt.tsx b/frontend/MASSter-frontend/src/components/video-cover-prompt.tsx
@@ -162,17 +162,13 @@ export function VideoCoverAdvancedPromptForm({
         negativePrompt: z.string().max(500).optional(),
         style: z.string().default("Нет"),
         inputImage: z
-            .custom<FileList>((v) => v instanceof FileList, {
+            .custom<File>((v) => v instanceof File, {
                 message: "Нужно загрузить картинки для референса"
             })
-            .refine((files) => {
-                for (let file of files) {
-                    if (!ACCEPTED_IMAGE_TYPES.includes(file.type)) {
-                        return false;
-                    }
-                }
-                return true;
-            }, "Пока мы поддерживаем только .png и .jpg изображения")
+            .refine(
+                (v) => ACCEPTED_IMAGE_TYPES.includes(v.type),
+                "Пока мы поддерживаем только .png и .jpg изображения"
+            )
             .optional(),
         video: z
             .custom<File>((v) => v instanceof File, {
@@ -204,9 +200,7 @@ export function VideoCoverAdvancedPromptForm({
         formData.append("num_images", `${data.numImages}`);
 
         if (data.inputImage) {
-            for (let file of data.inputImage) {
-                formData.append("input_image", file);
-            }
+            formData.append("input_image", data.inputImage);
         }
         if (data.video) {
             formData.append("video", data.video);
@@ -338,7 +332,6 @@ export function VideoCoverAdvancedPromptForm({
                                             <FormControl>
                                                 <Input
                                                     type="file"
-                                                    multiple={true}
                                                     onBlur={field.onBlur}
                                                     onChange={(e) =>
                                                         field.onChange(

diff --git a/frontend/MASSter-frontend/src/lib/pb-client.ts b/frontend/MASSter-frontend/src/lib/pb-client.ts
@@ -1,6 +1,6 @@
 import PocketBase from "pocketbase";
 
-export const pb = new PocketBase("https://pb.apps.npod.space");
+export const pb = new PocketBase(import.meta.env.VITE_PB_HOST);
 // pb.autoCancellation(false);
 
 
diff --git a/model/generator/libs/models.py b/model/generator/libs/models.py
@@ -59,6 +59,8 @@ def remove_models(self):
         except:
             try:
                 del self.pipe
+                del self.pipe_prior
+                del self.depth_estimator
             except:
                 ...
 
@@ -157,11 +159,7 @@ def generate_based_on_image(
                 +'low quality, jpeg artifacts, ugly, duplicate, '
                 + 'morbid, mutilated, out of frame, extra fingers, '
                 + 'mutated hands, poorly drawn hands, poorly drawn face, '
-                + 'mutation, deformed, blurry, dehydrated, bad anatomy, '
-                + 'bad proportions, extra limbs, cloned face, disfigured, '
-                + 'gross proportions, malformed limbs, missing arms, '
-                + 'missing legs, extra arms, extra legs, fused fingers, '
-                + 'too many fingers, long neck, username, watermark, signature'
+                + 'mutation, deformed, blurry, dehydrated, bad anatomy'
             )
 
         if num_steps < 1: num_steps = 1
@@ -184,7 +182,6 @@ def generate_based_on_image(
 
             images = self.pipe(
                 image_embeds=image_emb,
-                guidance_scale=guidance_scale,
                 negative_image_embeds=zero_image_emb,
                 hint=hint,
                 num_inference_steps=num_steps,

diff --git a/model/generator/main.py b/model/generator/main.py
@@ -90,25 +90,37 @@ def run(self):
     def process_generation(self, data, width, height) -> list[Image.Image]:
         # self.pb.get_file_url(data, filename=f'tmp/{data.record.image}')
 
-        if data.input_image != '':
+        if data.input_image != '' and data.type == 'avatar':
             try:
                 image_url = self.pb.collection('text_generation_mvp').get_file_url(data, filename=data.input_image) 
                 urllib.request.urlretrieve(image_url, data.input_image)
 
+                reference_image = Image.open(data.input_image, mode='r')
+                reference_image = reference_image.resize(size=(768, 768))
+                reference_image.save(data.input_image)
+
+                self.model.remove_models()
+                self.model.load_depth_model()
                 images = self.model.generate_based_on_image(
-                    image_path=image_url,
+                    image_path=data.input_image,
                     num_steps=50,
                     guidance_scale=4.0,
-                    height=height,
-                    width=width,
+                    height=768,
+                    width=768,
                     prompt=data.prompt,
                     style=data.style,
                     negative_prompt=data.negative_prompt,
                     num_images=data.num_images
                 )
+                self.model.remove_depth_model()
+                self.model.load_kandinsky()
+
+                for image in images:
+                    image = image.resize(size=(800, 800))
 
                 return images
             except:
+                raise
                 logger.error('Ошибка при скачивании референса, генерирую без референса')
 
             images = self.model.generate(

diff --git a/model/generator/research/funny_portrait_of_an_racoon_AH5I5YwSsv.jpg b/model/generator/research/funny_portrait_of_an_racoon_AH5I5YwSsv.jpg
diff --git a/model/generator/research/kandinsky.ipynb b/model/generator/research/kandinsky.ipynb
diff --git a/pocketbase/Dockerfile b/pocketbase/Dockerfile
@@ -0,0 +1,14 @@
+FROM debian:11
+
+WORKDIR /pocketbase
+
+RUN apt update -y
+RUN apt install -y unzip wget
+
+RUN wget https://github.com/pocketbase/pocketbase/releases/download/v0.18.10/pocketbase_0.18.10_linux_amd64.zip
+RUN unzip pocketbase_0.18.10_linux_amd64.zip
+
+COPY ./pb_migrations ./pb_migrations
+
+WORKDIR /
+