use new runwayml inpaint model 🎉

frontend/src/lib/App.svelte

@@ -71,12 +71,15 @@
 		const sessionHash = crypto.randomUUID();
 		const base64Crop = $maskEl.toDataURL('image/png');
 
-		const payload = {
+		const hashpayload = {
 			fn_index: 0,
-			data: [base64Crop, prompt, 0.75, 7.5, 35, 'patchmatch'],
 			session_hash: sessionHash
 		};
 
+		const datapayload = {
+			data: [base64Crop, prompt, 0.75, 7.5, 35, 'patchmatch']
+		}
+
 		const websocket = new WebSocket(PUBLIC_WS_INPAINTING);
 		// websocket.onopen = async function (event) {
 		// 	websocket.send(JSON.stringify({ hash: sessionHash }));
@@ -94,9 +97,12 @@
 				const data = JSON.parse(event.data);
 				$loadingState = '';
 				switch (data.msg) {
+					case 'send_hash':
+						websocket.send(JSON.stringify(hashpayload));
+						break;
 					case 'send_data':
 						$loadingState = 'Sending Data';
-						websocket.send(JSON.stringify(payload));
+						websocket.send(JSON.stringify({...hashpayload, ...datapayload}));
 						break;
 					case 'queue_full':
 						$loadingState = 'Queue full';

frontend/src/lib/Buttons/RoomsSelector.svelte

@@ -25,10 +25,10 @@
 	}
 	onMount(() => {
 		refreshRooms();
-		window.addEventListener('click', clickHandler, true);
+		window.addEventListener('pointerdown', clickHandler, true);
 		const interval = setInterval(refreshRooms, 3000);
 		return () => {
-			window.removeEventListener('click', clickHandler, true);
+			window.removeEventListener('pointerdown', clickHandler, true);
 			clearInterval(interval);
 		};
 	});

frontend/src/lib/PromptModal.svelte

@@ -23,11 +23,12 @@
 		inputEl.focus();
 		prompt = initPrompt;
 		window.addEventListener('keyup', onKeyup);
-		window.addEventListener('click', cancelHandler, true);
+		window.addEventListener('pointerdown', cancelHandler, true);
+
 
 		return () => {
 			window.removeEventListener('keyup', onKeyup);
-			window.removeEventListener('click', cancelHandler, true);
+			window.removeEventListener('pointerdown', cancelHandler, true);
 		};
 	});
 	let timer: NodeJS.Timeout;

requirements.txt

@@ -1,15 +1,15 @@
 --extra-index-url https://download.pytorch.org/whl/cu113 
 torch
 huggingface_hub
-diffusers
-transformers
-scikit-image
-pillow
-opencv-python-headless
-fastapi
-uvicorn
-httpx
-gradio
-boto3
-python-magic
-fastapi-utils
+diffusers==0.6.0
+transformers==4.23.1
+scikit-image==0.19.3
+Pillow==9.2.0
+opencv-python-headless==4.6.0.66
+fastapi==0.85.1
+uvicorn==0.18.3
+httpx==0.23.0
+gradio==3.6
+boto3==1.24.93
+python-magic==0.4.27
+fastapi-utils==0.2.1

stablediffusion-infinity/app.py

@@ -71,21 +71,11 @@ model = {}
 
 
 def get_model():
-    if "text2img" not in model:
-        text2img = StableDiffusionPipeline.from_pretrained(
-            "CompVis/stable-diffusion-v1-4",
+    if "inpaint" not in model:
+        inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
             revision="fp16",
-            torch_dtype=torch.float16,
-            use_auth_token=HF_TOKEN,
-        ).to("cuda")
-        inpaint = StableDiffusionInpaintPipeline(
-            vae=text2img.vae,
-            text_encoder=text2img.text_encoder,
-            tokenizer=text2img.tokenizer,
-            unet=text2img.unet,
-            scheduler=text2img.scheduler,
-            safety_checker=text2img.safety_checker,
-            feature_extractor=text2img.feature_extractor,
+            torch_dtype=torch.float16
         ).to("cuda")
 
         # lms = LMSDiscreteScheduler(
@@ -108,11 +98,10 @@ def get_model():
         #         inpaint.enable_attention_slicing()
         # except:
         #     pass
-        model["text2img"] = text2img
         model["inpaint"] = inpaint
         # model["img2img"] = img2img
 
-    return model["text2img"], model["inpaint"]
+    return model["inpaint"]
     # model["img2img"]
 
 
@@ -127,8 +116,10 @@ def run_outpaint(
     guidance,
     step,
     fill_mode,
+
+
 ):
-    text2img, inpaint = get_model()
+    inpaint = get_model()
     sel_buffer = np.array(input_image)
     img = sel_buffer[:, :, 0:3]
     mask = sel_buffer[:, :, -1]
@@ -167,25 +158,30 @@ def run_outpaint(
         mask_image = Image.fromarray(mask)
 
         # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
-        with autocast("cuda"):
-            images = inpaint(
-                prompt=prompt_text,
-                init_image=init_image.resize(
-                    (process_size, process_size), resample=SAMPLING_MODE
-                ),
-                mask_image=mask_image.resize((process_size, process_size)),
-                strength=strength,
-                num_inference_steps=step,
-                guidance_scale=guidance,
-            )
+
     else:
         print("text2image")
-        with autocast("cuda"):
-            images = text2img(
-                prompt=prompt_text, height=process_size, width=process_size,
-            )
+        print("inpainting")
+        img, mask = functbl[fill_mode](img, mask)
+        init_image = Image.fromarray(img)
+        mask = 255 - mask
+        mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+        mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+        mask_image = Image.fromarray(mask)
 
-    return images['sample'][0], images["nsfw_content_detected"][0]
+        # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+    with autocast("cuda"):
+        output = inpaint(
+            prompt=prompt_text,
+            image=init_image.resize(
+                (process_size, process_size), resample=SAMPLING_MODE
+            ),
+            mask_image=mask_image.resize((process_size, process_size)),
+            strength=strength,
+            num_inference_steps=step,
+            guidance_scale=guidance,
+        )
+    return output['images'][0], output["nsfw_content_detected"][0]
 
 
 with blocks as demo:
@@ -242,6 +238,8 @@ blocks.config['dev_mode'] = False
 
 app = gr.mount_gradio_app(app, blocks, "/gradio",
                           gradio_api_url="http://0.0.0.0:7860/gradio/")
+
+
 def generateAuthToken():
     response = requests.get(f"https://liveblocks.io/api/authorize",
                             headers={"Authorization": f"Bearer {LIVEBLOCKS_SECRET}"})