|
|
|
|
|
from diffusers import StableDiffusionInpaintPipeline |
|
|
import requests |
|
|
import torch |
|
|
from PIL import Image |
|
|
from io import BytesIO |
|
|
|
|
|
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" |
|
|
|
|
|
def download_image(url): |
|
|
response = requests.get(url) |
|
|
return 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" |
|
|
|
|
|
init_image = download_image(img_url).resize((512, 512)) |
|
|
mask_image = download_image(mask_url).resize((512, 512)) |
|
|
|
|
|
path = "runwayml/stable-diffusion-inpainting" |
|
|
|
|
|
run_compile = True |
|
|
|
|
|
pipe = StableDiffusionInpaintPipeline.from_pretrained(path, torch_dtype=torch.float16) |
|
|
pipe = pipe.to("cuda:0") |
|
|
pipe.unet.to(memory_format=torch.channels_last) |
|
|
|
|
|
if run_compile: |
|
|
print("Run torch compile") |
|
|
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) |
|
|
|
|
|
prompt = "ghibli style, a fantasy landscape with castles" |
|
|
|
|
|
for _ in range(3): |
|
|
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0] |
|
|
|
