pip install torch numpy torchvision diffusers

import torch as th
import numpy as np
import torchvision.utils as tvu
from diffusers import DiffusionPipeline
import argparse
import sys

# Simulate passing arguments explicitly (bypassing Jupyter's arguments)
sys.argv = [
    "ipykernel_launcher.py", 
    "--prompt", "mystical trees | A magical pond | dark",
    "--steps", "50",
    "--scale", "7.5",
    "--weights", "7.5 | 7.5 | -7.5",
    "--seed", "2",
    "--model_path", "CompVis/stable-diffusion-v1-4",
    "--num_images", "1"
]

parser = argparse.ArgumentParser()
parser.add_argument("--prompt", type=str, default="mystical trees | A magical pond | dark",
                    help="use '|' as the delimiter to compose separate sentences.")
parser.add_argument("--steps", type=int, default=50)
parser.add_argument("--scale", type=float, default=7.5)
parser.add_argument("--weights", type=str, default="7.5 | 7.5 | -7.5")
parser.add_argument("--seed", type=int, default=2)
parser.add_argument("--model_path", type=str, default="CompVis/stable-diffusion-v1-4")
parser.add_argument("--num_images", type=int, default=1)
args = parser.parse_args()

# CUDA Setup
has_cuda = th.cuda.is_available()
device = th.device('cpu' if not has_cuda else 'cuda')

# Assign parameters
prompt = args.prompt
scale = args.scale
steps = args.steps

# Load pipeline
pipe = DiffusionPipeline.from_pretrained(
    args.model_path,
    custom_pipeline="composable_stable_diffusion",
).to(device)

# Disable safety checker (if intentional for internal use)
pipe.safety_checker = None

# Generate images
images = []
generator = th.Generator("cuda").manual_seed(args.seed)
for i in range(args.num_images):
    image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps,
                 weights=args.weights, generator=generator).images[0]
    images.append(th.from_numpy(np.array(image)).permute(2, 0, 1) / 255.)

# Create and save image grid
grid = tvu.make_grid(th.stack(images, dim=0), nrow=4, padding=0)
tvu.save_image(grid, f'{prompt}_{args.weights}.png')

print("Image saved successfully!")