├── cog.yaml
├── script
    └── download-weights
├── README.md
└── predict.py


/cog.yaml:
--------------------------------------------------------------------------------
 1 | build:
 2 |   gpu: true
 3 |   python_version: "3.10"
 4 |   python_packages:
 5 |     - "diffusers==0.9.0"
 6 |     - "torch==1.13.0"
 7 |     - "ftfy==6.1.1"
 8 |     - "scipy==1.9.3"
 9 |     - "transformers==4.25.1"
10 |     - "accelerate==0.15.0"
11 | 
12 | predict: "predict.py:Predictor"
13 | 


--------------------------------------------------------------------------------
/script/download-weights:
--------------------------------------------------------------------------------
 1 | #!/usr/bin/env python
 2 | 
 3 | import os
 4 | from diffusers import StableDiffusionInpaintPipeline
 5 | 
 6 | 
 7 | model_id = "stabilityai/stable-diffusion-2-inpainting"
 8 | MODEL_CACHE = "diffusers-cache"
 9 | os.makedirs(MODEL_CACHE, exist_ok=True)
10 | 
11 | pipe = StableDiffusionInpaintPipeline.from_pretrained(
12 |     model_id,
13 |     cache_dir=MODEL_CACHE,
14 | )
15 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # cog-stable-diffusion-inpainting-v2
 2 | 
 3 | 
 4 | [![Replicate](https://replicate.com/replicate/stable-diffusion-inpainting/badge)](https://replicate.com/replicate/stable-diffusion-inpainting) 
 5 | 
 6 | This is an implementation of the [Diffusers Stable Diffusion v2](https://huggingface.co/stabilityai/stable-diffusion-2) as a Cog model. [Cog packages machine learning models as standard containers.](https://github.com/replicate/cog)
 7 | 
 8 | First, download the pre-trained weights:
 9 | 
10 |     cog run script/download-weights 
11 | 
12 | Then, you can run predictions:
13 | 
14 |     cog predict -i prompt="..." -i image=@... -i mask=@...
15 | 


--------------------------------------------------------------------------------
/predict.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | from typing import List
 3 | 
 4 | import torch
 5 | from diffusers import StableDiffusionInpaintPipeline
 6 | 
 7 | from PIL import Image
 8 | from cog import BasePredictor, Input, Path
 9 | 
10 | MODEL_ID = "stabilityai/stable-diffusion-2-inpainting"
11 | MODEL_CACHE = "diffusers-cache"
12 | 
13 | 
14 | class Predictor(BasePredictor):
15 |     def setup(self):
16 |         """Load the model into memory to make running multiple predictions efficient"""
17 |         print("Loading pipeline...")
18 |         self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
19 |             MODEL_ID,
20 |             cache_dir=MODEL_CACHE,
21 |             local_files_only=True,
22 |         ).to("cuda")
23 | 
24 |     @torch.inference_mode()
25 |     @torch.cuda.amp.autocast()
26 |     def predict(
27 |         self,
28 |         prompt: str = Input(
29 |             description="Input prompt",
30 |             default="a photo of an astronaut riding a horse on mars",
31 |         ),
32 |         negative_prompt: str = Input(
33 |             description="Specify things to not see in the output",
34 |             default=None
35 |         ),
36 |         image: Path = Input(
37 |             description="Inital image to generate variations of. Supproting images size with 512x512",
38 |         ),
39 |         mask: Path = Input(
40 |             description="Black and white image to use as mask for inpainting over the image provided. White pixels are inpainted and black pixels are preserved",
41 |         ),
42 |         num_outputs: int = Input(
43 |             description="Number of images to output. Higher number of outputs may OOM.",
44 |             ge=1,
45 |             le=8,
46 |             default=1,
47 |         ),
48 |         num_inference_steps: int = Input(
49 |             description="Number of denoising steps", ge=1, le=500, default=25
50 |         ),
51 |         guidance_scale: float = Input(
52 |             description="Scale for classifier-free guidance", ge=1, le=20, default=7.5
53 |         ),
54 |         seed: int = Input(
55 |             description="Random seed. Leave blank to randomize the seed", default=None
56 |         ),
57 |     ) -> List[Path]:
58 |         """Run a single prediction on the model"""
59 |         if seed is None:
60 |             seed = int.from_bytes(os.urandom(2), "big")
61 |         print(f"Using seed: {seed}")
62 | 
63 |         image = Image.open(image).convert("RGB").resize((512, 512))
64 |         extra_kwargs = {
65 |             "mask_image": Image.open(mask).convert("RGB").resize(image.size),
66 |             "image": image
67 |         }
68 | 
69 |         generator = torch.Generator("cuda").manual_seed(seed)
70 | 
71 |         output = self.pipe(
72 |             prompt=[prompt] * num_outputs if prompt is not None else None,
73 |             negative_prompt=[negative_prompt] * num_outputs
74 |             if negative_prompt is not None
75 |             else None,
76 |             guidance_scale=guidance_scale,
77 |             generator=generator,
78 |             num_inference_steps=num_inference_steps,
79 |             **extra_kwargs,
80 |         )
81 | 
82 |         output_paths = []
83 |         for i, sample in enumerate(output.images):
84 |             output_path = f"/tmp/out-{i}.png"
85 |             sample.save(output_path)
86 |             output_paths.append(Path(output_path))
87 | 
88 |         return output_paths
89 | 


--------------------------------------------------------------------------------