rearranged fun

app.py

@@ -9,34 +9,65 @@ from rembg import remove
 import requests
 from io import BytesIO
 
-# def image_grid(imgs, rows, cols):
-#     assert len(imgs) == rows*cols
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows*cols
 
-#     w, h = imgs[0].size
-#     grid = PIL.Image.new('RGB', size=(cols*w, rows*h))
-#     grid_w, grid_h = grid.size
+    w, h = imgs[0].size
+    grid = PIL.Image.new('RGB', size=(cols*w, rows*h))
+    grid_w, grid_h = grid.size
+    
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid
     
-#     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 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")
 
-def greet(name):
-    img_url = "https://cdn.faire.com/fastly/893b071985d70819da5f0d485f1b1bb97ee4f16a6e14ef1bdd4a086b3588be58.png" # wino
-    image = download_image(img_url).resize((512, 512))
-    inverted_mask_image = remove(data = image, only_mask = True)
-    mask_image = PIL.ImageOps.invert(inverted_mask_image)
+device = "cuda"
+model_path = "runwayml/stable-diffusion-inpainting"
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    model_path,
+    revision="fp16", 
+    torch_dtype=torch.float16,
+    use_auth_token=True
+).to(device)
+
+img_url = "https://cdn.faire.com/fastly/893b071985d70819da5f0d485f1b1bb97ee4f16a6e14ef1bdd4a086b3588be58.png" # wino
+inverted_mask_image = remove(data = image, only_mask = True)
+mask_image = PIL.ImageOps.invert(inverted_mask_image)
+prompt = "crazy portal universe"
 
-    return "Hello " + name + "!!"
+guidance_scale=7.5
+num_samples = 3
+generator = torch.Generator(device="cuda").manual_seed(0) # change the seed to get different results
+images = pipe(
+    prompt=prompt,
+    image=image,
+    mask_image=mask_image,
+    guidance_scale=guidance_scale,
+    generator=generator,
+    num_images_per_prompt=num_samples,
+).images
+images.insert(0, image)
+image_grid(images, 1, num_samples + 1)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+gr.Interface(
+    predict,
+    title = 'Stable Diffusion In-Painting',
+    inputs=[
+        gr.Image(source = 'upload', tool = 'sketch', type = 'pil'),
+        gr.Textbox(label = 'prompt')
+    ],
+    outputs = [
+        gr.Image()
+        ]
+).launch(debug=True)