Update README.md

README.md

@@ -5,8 +5,21 @@ base_model:
 pipeline_tag: image-to-image
 ---
 ```
+import torch
+import numpy as np
+
+from PIL import Image
+from diffusers.utils import load_image
 from transformers import AutoConfig, AutoModel, ModelCard
 
+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"
+
+image_gray = load_image(img_url).resize((512, 512)).convert('L').convert('RGB') # image must be 3 channel
+mask_image = load_image(mask_url).resize((512, 512))
+mask = (np.array(mask_image)>128)*1
+image_gray_masked = Image.fromarray(((1-mask) * np.array(image_gray)).astype(np.uint8))
+
 # Load the gray-inpaint model
 gray_inpaintor = AutoModel.from_pretrained(
     'jwengr/stable-diffusion-2-gray-inpaint-to-rgb', 
@@ -14,25 +27,24 @@ gray_inpaintor = AutoModel.from_pretrained(
     trust_remote_code=True, 
 )
 
-# Load the gray2rgb model
+Load the gray2rgb model
 gray2rgb = AutoModel.from_pretrained(
     'jwengr/stable-diffusion-2-gray-inpaint-to-rgb', 
     subfolder='gray2rgb', 
     trust_remote_code=True, 
 )
 
-
-# Move models to GPU
-gray2rgb.to('cuda')
+Move models to GPU
 gray_inpaintor.to('cuda')
+gray2rgb.to('cuda')
 
-# Enable memory-efficient attention for xFormers
-gray2rgb.unet.enable_xformers_memory_efficient_attention()
-gray_inpaintor.unet.enable_xformers_memory_efficient_attention()
+# Enable memory-efficient attention
+# gray2rgb.unet.enable_xformers_memory_efficient_attention()
+# gray_inpaintor.unet.enable_xformers_memory_efficient_attention()
 
-# Generate images using gray_inpaintor and gray2rgb
-image_gray_restored = gray_inpaintor(batch, seed=inpaint_seed)
-image_gray_restored = [img.convert('RGB') for img in image_gray_restored]
-image_restored_pil = gray2rgb(image_gray_restored)
-image_restored_pt = gray2rgb(image_gray_restored, output_type='pt')
+with torch.autocast('cuda',dtype=torch.bfloat16):
+    with torch.no_grad():
+        # each model's input image should be one of PIL.Image, List[PIL.Image], preprocessed tensor (B,3,H,W). Image must be 3 channel
+        image_gray_restored = gray_inpaintor(image_gray_masked, num_inference_steps=250, seed=10)[0].convert('L') # you can pass 'mask' arg explictly. mask : Tensor (B,1,512,512)
+        image_restored = gray2rgb(image_gray_restored.convert('RGB'))
 ```