Spaces:
Sleeping
Sleeping
File size: 3,304 Bytes
43c5517 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from diffusers import StableDiffusionInpaintPipeline
import torch
from PIL import Image
import os
from torch.multiprocessing import set_start_method
import gc
def setup_model(model_path):
# Clear CUDA memory
torch.cuda.empty_cache()
gc.collect()
pipe = StableDiffusionInpaintPipeline.from_single_file(
model_path,
torch_dtype=torch.float16,
safety_checker=None
).to("cuda")
# Enable memory optimizations without xformers
pipe.enable_attention_slicing(slice_size="max")
pipe.enable_sequential_cpu_offload()
return pipe
def prepare_images(image_path, mask_path=None):
# Load and prepare the original image
original_image = Image.open(image_path)
# Resize to a multiple of 8 (required by Stable Diffusion)
width, height = (dim - dim % 8 for dim in original_image.size)
original_image = original_image.resize((width, height))
if mask_path:
mask_image = Image.open(mask_path)
mask_image = mask_image.resize((width, height))
mask_image = mask_image.convert("L")
else:
# Create a simple rectangular mask in the center
mask_image = Image.new("L", (width, height), 0)
mask_width = width // 3
mask_height = height // 3
x1 = (width - mask_width) // 2
y1 = (height - mask_height) // 2
x2 = x1 + mask_width
y2 = y1 + mask_height
for y in range(y1, y2):
for x in range(x1, x2):
mask_image.putpixel((x, y), 255)
return original_image, mask_image
def main():
# Setup paths using raw strings
model_path = "realisticVisionV60B1_v51VAE-inpainting.safetensors"
image_path = r"C:\Users\M. Y\Downloads\t2.png"
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name()}")
print(f"Memory allocated: {torch.cuda.memory_allocated()/1024**2:.2f}MB")
# Initialize model
print("Loading model...")
pipe = setup_model(model_path)
# Prepare images
print("Preparing images...")
original_image, mask_image = prepare_images(image_path)
# Save mask for verification
mask_image.save("generated_mask.png")
mask_image_1 = Image.open("generated_mask_1.png")
# Define your prompt
prompt = "add some flowers and a fountain"
negative_prompt = "blurry, low quality, distorted"
print("Performing inpainting...")
with torch.cuda.amp.autocast(): # Use automatic mixed precision
output_image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image=original_image,
mask_image=mask_image_1,
num_inference_steps=20, # Reduced steps for faster generation
guidance_scale=7.5,
).images[0]
# Save the result
output_image.save("inpainted_result.png")
print("Inpainting completed! Check 'inpainted_result.png' for the result.")
# Clean up
torch.cuda.empty_cache()
gc.collect()
if __name__ == "__main__":
try:
set_start_method('spawn')
except RuntimeError:
pass
main() |