import os
import torch
from diffusers import StableDiffusion3Img2ImgPipeline
from diffusers.utils import load_image
from PIL import Image

# Set device and model path
device = "cuda"
model_id_or_path = "stabilityai/stable-diffusion-3-medium-diffusers"

# Initialize pipeline
pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
pipe.load_lora_weights("zwloong/sd3-lora-training-rank16-v2")
pipe = pipe.to(device)

# Load the initial image
init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png")

# Set prompt and negative prompt
prompt = "A design of a cute pokemon, featuring a single bird, water type pokemon, on a background"
negative_prompt = 'blurry, cropped, ugly'

# List of different scale values to apply
scale_values = [0.5, 0.7, 0.9, 1.1, 1.3]

# Directory to save images
save_dir = "./generated_images_scale"
os.makedirs(save_dir, exist_ok=True)

# Loop over each scale value
for i, scale in enumerate(scale_values):
    # Generate the image with the current scale value
    generator = torch.Generator(device='cuda').manual_seed(44)
    image = pipe(
        prompt=prompt,
        image=init_image,
        generator=generator,
        guidance_scale=5.0,
        negative_prompt=negative_prompt,
        joint_attention_kwargs={"scale": scale}
    ).images[0]
    
    # Save the image
    image_save_path = os.path.join(save_dir, f"image_scale_{scale}.png")
    image.save(image_save_path)
    print(f"Saved image with scale {scale} to {image_save_path}")