ImgToImgSusFast.py

from pathlib import Path
import re
import torch
from PIL import Image, ImageOps
from diffusers import StableDiffusionXLImg2ImgPipeline
import tkinter as tk
from tkinter import filedialog

model_id = "RunDiffusion/Juggernaut-XL-v9"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

def safe_filename(text: str, max_len: int = 80) -> str:
    text = text.strip()
    text = re.sub(r"[^\w\s.-]", "", text, flags=re.UNICODE)
    text = re.sub(r"\s+", "_", text)
    text = text.strip("._-")
    return (text or "edited")[:max_len]

print("Loading Juggernaut XL v9...")
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    model_id, torch_dtype=dtype, use_safetensors=True, variant="fp16"
)
pipe = pipe.to(device)
pipe.enable_attention_slicing()
pipe.vae.enable_slicing()

root = tk.Tk()
root.withdraw()

print("Select input image...")
file_path = filedialog.askopenfilename(title="Select Image", filetypes=[("Image files", "*.png *.jpg *.jpeg *.webp *.bmp")])
if not file_path: raise SystemExit("No file selected.")

image_path = Path(file_path)

prompt = input("\nEnter your prompt: ").strip()
if not prompt: raise SystemExit("No prompt entered.")

negative_prompt = input("Negative prompt (Enter for default): ").strip() or "blurry, deformed, ugly, bad anatomy, extra limbs, low quality, worst quality, text, watermark"

with Image.open(image_path) as img:
    init_image = ImageOps.exif_transpose(img).convert("RGB")
    init_image.thumbnail((512, 512), Image.Resampling.LANCZOS)

safe_prompt = safe_filename(prompt)
output_path = image_path.with_name(f"{image_path.stem}_{safe_prompt}.png")

print("Generating image...")
with torch.inference_mode():
    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=init_image,
        strength=0.55,
        guidance_scale=5.5,
        num_inference_steps=8
    )

result.images[0].save(output_path)
print(f"\nImage saved to: {output_path.resolve()}")