File size: 1,257 Bytes
52722af
255d038
 
 
 
 
52722af
255d038
52722af
e7b2bab
52722af
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
96be5bf
52722af
 
 
 
 
 
 
 
6db47b5
44e42ee
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import streamlit as st
import torch
import numpy
from PIL import Image
from torchvision import transforms
from diffusers import StableDiffusionInpaintPipeline
from diffusers import DPMSolverMultistepScheduler, UniPCMultistepScheduler

@torch.inference_mode()
@st.cache_resource
def get_pipeline():
    pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting",
                                                                      torch_dtype=torch.float16)
    pipe.to(device)
    pipe.enable_xformers_memory_efficient_attention()
    pipe.set_progress_bar_config(disable=True)
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
    return pipe


def inpainting(image,
               mask_image,
               prompt,
               negative_prompt,
               num_inference_steps=20,
               guidance_scale=7.5,
               ):
    pipe = get_pipeline()
    print("retrieved pipeline")
    result = pipe(
        image=image,
        mask_image=mask_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
    ).images[0]
    print("Generated image")
    return result