Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,316 Bytes
bd199cf |
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 |
import torch
from PIL import Image
from torchvision import transforms
from diff_pipe import StableDiffusionXLDiffImg2ImgPipeline
device = "cuda"
base = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to(device)
refiner = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0",
text_encoder_2=base.text_encoder_2,
vae=base.vae,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
).to(device)
def preprocess_image(image):
image = image.convert("RGB")
image = transforms.CenterCrop((image.size[1] // 64 * 64, image.size[0] // 64 * 64))(image)
image = transforms.ToTensor()(image)
image = image * 2 - 1
image = image.unsqueeze(0).to(device)
return image
def preprocess_map(map):
map = map.convert("L")
map = transforms.CenterCrop((map.size[1] // 64 * 64, map.size[0] // 64 * 64))(map)
# convert to tensor
map = transforms.ToTensor()(map)
map = map.to(device)
return map
with Image.open("assets/input2.jpg") as imageFile:
image = preprocess_image(imageFile)
with Image.open("assets/map2.jpg") as mapFile:
map = preprocess_map(mapFile)
prompt = ["painting of a mountain landscape with a meadow and a forest, meadow background"]
negative_prompt = ["blurry, shadow polaroid photo, scary angry pose"]
edited_images = base(prompt=prompt, original_image=image, image=image, strength=1, guidance_scale=17.5,
num_images_per_prompt=1,
negative_prompt=negative_prompt,
map=map,
num_inference_steps=100, denoising_end=0.8, output_type="latent").images
edited_images = refiner(prompt=prompt, original_image=image, image=edited_images, strength=1, guidance_scale=17.5,
num_images_per_prompt=1,
negative_prompt=negative_prompt,
map=map,
num_inference_steps=100, denoising_start=0.8).images[0]
# Despite we use here both of the refiner and the base models,
# one can use only the base model, or only the refiner (for low strengths).
edited_images.save("output.png")
print("Done!")
|