File size: 756 Bytes
b8877ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | import torch
from diffusers import StableDiffusionPipeline, DiffusionPipeline
def test_sdg_minimal():
sdg = DiffusionPipeline.from_pretrained(
"your-org/safe-diffusion-guidance",
custom_pipeline="safe_diffusion_guidance",
torch_dtype=torch.float16
)
sdg = sdg.to("cuda" if torch.cuda.is_available() else "cpu")
out = sdg(
prompt="test scene",
base_model_id="runwayml/stable-diffusion-v1-5",
num_inference_steps=2, # small for CI
guidance_scale=5.0,
safety_scale=2.0,
mid_fraction=0.5,
safe_class_index=3
)
assert len(out.images) == 1
print("OK: pipeline end-to-end")
if __name__ == "__main__":
test_sdg_minimal()
|