Update app.py

app.py

@@ -1,24 +1,18 @@
-from huggingface_hub import notebook_login
-
-notebook_login()
 import inspect
+import os
 from typing import List, Optional, Union
-
 import numpy as np
 import torch
-
 import PIL
 import gradio as gr
 from diffusers import StableDiffusionInpaintPipeline
-device = "cuda"
-model_path = "runwayml/stable-diffusion-inpainting"
-
-pipe = StableDiffusionInpaintPipeline.from_pretrained(
-    model_path,
-    torch_dtype=torch.float16,
-).to(device)
+from rembg import remove
 import requests
 from io import BytesIO
+from huggingface_hub import login
+
+token = os.getenv("WRITE_TOKEN")
+login(token, True)
 
 def image_grid(imgs, rows, cols):
     assert len(imgs) == rows*cols
@@ -30,24 +24,36 @@ def image_grid(imgs, rows, cols):
     for i, img in enumerate(imgs):
         grid.paste(img, box=(i%cols*w, i//cols*h))
     return grid
+    
 
+def predict(dict, prompt):
+  image = dict['image'].convert("RGB").resize((512, 512))
+  mask_image = dict['mask'].convert("RGB").resize((512, 512))
+  images = pipe(prompt=prompt, image=image, mask_image=mask_image).images
+  return(images[0])
 
 def download_image(url):
     response = requests.get(url)
     return PIL.Image.open(BytesIO(response.content)).convert("RGB")
 
-img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
-mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+model_path = "runwayml/stable-diffusion-inpainting"
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    model_path,
+    # revision="fp16", 
+    # torch_dtype=torch.float16,
+    use_auth_token=True
+)
+
+img_url = "https://cdn.faire.com/fastly/893b071985d70819da5f0d485f1b1bb97ee4f16a6e14ef1bdd4a086b3588be58.png" # wino
 image = download_image(img_url).resize((512, 512))
-image
-mask_image = download_image(mask_url).resize((512, 512))
-mask_image
-prompt = "a mecha robot sitting on a bench"
+inverted_mask_image = remove(data = image, only_mask = True)
+mask_image = PIL.ImageOps.invert(inverted_mask_image)
+prompt = "crazy portal universe"
 
 guidance_scale=7.5
 num_samples = 3
-generator = torch.Generator(device="cuda").manual_seed(0) # change the seed to get different results
-
+generator = torch.Generator(device="cpu").manual_seed(0) # change the seed to get different results
 images = pipe(
     prompt=prompt,
     image=image,
@@ -56,15 +62,10 @@ images = pipe(
     generator=generator,
     num_images_per_prompt=num_samples,
 ).images
-# insert initial image in the list so we can compare side by side
 images.insert(0, image)
 image_grid(images, 1, num_samples + 1)
-def predict(dict, prompt):
-  image =  dict['image'].convert("RGB").resize((512, 512))
-  mask_image = dict['mask'].convert("RGB").resize((512, 512))
-  images = pipe(prompt=prompt, image=image, mask_image=mask_image).images
-  return(images[0])
-  gr.Interface(
+
+gr.Interface(
     predict,
     title = 'Stable Diffusion In-Painting',
     inputs=[
@@ -74,4 +75,4 @@ def predict(dict, prompt):
     outputs = [
         gr.Image()
         ]
-).launch(debug=True)
+).launch(debug=True)