File size: 3,443 Bytes
7db9ab2
 
39544af
 
7db9ab2
 
39544af
 
 
 
 
 
 
 
 
 
8ca38ac
 
78392d4
39544af
 
78392d4
 
39544af
78392d4
 
39544af
78392d4
 
7db9ab2
78392d4
 
39544af
78392d4
 
 
 
 
 
 
 
 
 
 
 
39544af
 
7db9ab2
 
 
 
39544af
78392d4
39544af
 
78392d4
 
39544af
7db9ab2
5e772ca
78392d4
7db9ab2
 
f93e041
78392d4
7db9ab2
39544af
 
78392d4
 
 
7db9ab2
78392d4
39544af
7db9ab2
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import gradio
import torch
import numpy
from PIL import Image
from torchvision import transforms
from diffusers import StableDiffusionInpaintPipeline
from diffusers import DPMSolverMultistepScheduler

deviceStr = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(deviceStr)

if deviceStr == "cuda":
    pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting",
                                                              revision="fp16",
                                                              torch_dtype=torch.float16,
                                                              safety_checker=lambda images, **kwargs: (images, False))
    pipeline.to(device)
    pipeline.enable_xformers_memory_efficient_attention()
    latents = torch.randn((1, 4, 64, 64), device=device, dtype=torch.float16)
else:
    pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting",
                                                              safety_checker=lambda images, **kwargs: (images, False))                                                       
    latents = torch.randn((1, 4, 64, 64), device=device)

imageSize = (512, 512)
lastImage = Image.new(mode="RGB", size=imageSize)

lastSeed = 512
generator = torch.Generator(device).manual_seed(512)

def diffuse(staticLatents, inputImage, mask, pauseInference, prompt, negativePrompt, guidanceScale, numInferenceSteps, seed):
    global latents, lastSeed, generator, deviceStr, lastImage

    if mask is None or pauseInference is True:
        return lastImage
    
    if staticLatents is False:
        if deviceStr == "cuda":
            latents = torch.randn((1, 4, 64, 64), device=device, dtype=torch.float16)
        else:
            latents = torch.randn((1, 4, 64, 64), device=device)
    
    if lastSeed != seed:
        generator = torch.Generator(device).manual_seed(seed)
        lastSeed = seed
        
    newImage = pipeline(prompt=prompt,
                    negative_prompt=negativePrompt,
                    image=inputImage,
                    mask_image=mask,
                    guidance_scale=guidanceScale,
                    num_inference_steps=numInferenceSteps,
                    latents=latents,
                    generator=generator).images[0]

    lastImage = newImage

    return newImage

defaultMask = Image.open("assets/masks/sphere.png")

prompt = gradio.Textbox(label="Prompt", placeholder="A person in a room", lines=3)
negativePrompt = gradio.Textbox(label="Negative Prompt", placeholder="Text", lines=3)
inputImage = gradio.Image(label="Input Feed", source="webcam", shape=[512,512], streaming=True)
mask = gradio.Image(label="Mask", type="pil", value=defaultMask)
outputImage = gradio.Image(label="Extrapolated Field of View")
guidanceScale = gradio.Slider(label="Guidance Scale", maximum=1, value=0.75)
numInferenceSteps = gradio.Slider(label="Number of Inference Steps", maximum=100, value=25)
seed = gradio.Slider(label="Generator Seed", maximum=10000, value=4096)
staticLatents =gradio.Checkbox(label="Static Latents", value=True)
pauseInference = gradio.Checkbox(label="Pause Inference", value=False)

inputs=[staticLatents, inputImage, mask, pauseInference, prompt, negativePrompt, guidanceScale, numInferenceSteps, seed]
ux = gradio.Interface(fn=diffuse, title="View Diffusion", inputs=inputs, outputs=outputImage, live=True)
ux.launch()