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()