StableDiffusionImg2ImgPipeline

gistfile1.txt

import torch
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda" if torch.cuda.is_available() else "cpu"

pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(model_id)
pipeline = pipeline.to(device)

input_image_path = "image.png"
try:
    input_image = Image.open(input_image_path)
    input_image = input_image.convert("RGB")
except Exception as e:
    raise ValueError(f"Error loading imagen: {e}")

prompt = (
    "only modify the hairstyle of this person to add a futuristic haircut, "
    "without changing any other facial features or background."
)

init_image = input_image.resize((512, 512))

try:
    output = pipeline(
        prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5
    ).images[0]
except Exception as e:
    raise ValueError(f"Error generating imagen: {e}")

output_image_path = "output_image_modified.png"
output.save(output_image_path)

print(f"DONE... '{output_image_path}'")