Spaces:
Sleeping
Sleeping
from diffusers import StableDiffusionInpaintPipeline | |
import torch | |
from PIL import Image | |
def test_local(): | |
# Load model | |
model_id = "Uminosachi/realisticVisionV51_v51VAE-inpainting" | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
dtype = torch.float16 if device == "cuda" else torch.float32 | |
print(f"Using device: {device}") | |
pipe = StableDiffusionInpaintPipeline.from_pretrained( | |
model_id, | |
torch_dtype=dtype, | |
safety_checker=None | |
).to(device) | |
# Load test images | |
image_path = r"C:\Users\M. Y\Downloads\t2.png" | |
mask_path = "generated_mask_1.png" | |
image = Image.open(image_path) | |
mask_image = Image.open(mask_path) | |
# Resize to multiple of 8 | |
width, height = (dim - dim % 8 for dim in image.size) | |
image = image.resize((width, height)) | |
mask_image = mask_image.resize((width, height)) | |
mask_image = mask_image.convert("L") | |
# Test inference | |
result = pipe( | |
prompt="add some flowers and a fountain", | |
image=image, | |
mask_image=mask_image, | |
num_inference_steps=20, | |
guidance_scale=7.5, | |
).images[0] | |
result.save("local_test_result.png") | |
print("Test complete! Check local_test_result.png") | |
if __name__ == "__main__": | |
test_local() |