diffusers/tools

1#!/usr/bin/env python32import tree_ring_watermark as trk3from diffusers import DiffusionPipeline, DDIMScheduler4from pathlib import Path5from huggingface_hub import HfApi, login6import torch7 8 9# login()  # make sure you login it with on account that is connected to `trk-demo`10trk.set_org("trk-demo")11 12model_id = 'stabilityai/stable-diffusion-2-1-base'13device = 'cuda' if torch.cuda.is_available() else 'cpu'14 15# note that the model hash should be the latest commit hash of the repo's history: https://huggingface.co/stabilityai/stable-diffusion-2-base/commits/main16model_hash = "dcd3ee64f0c1aba2eb9e0c0c16041c6cae40d780"17 18pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)19pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)20pipe = pipe.to(device)21 22# get noise23batch_size = 124n_channels = pipe.unet.config.in_channels25sample_size = pipe.unet.config.sample_size 26 27shape = (batch_size, n_channels, sample_size, sample_size)28 29# get model hash from https://huggingface.co/stabilityai/stable-diffusion-2-1-base/commits/main30latents = trk.get_noise(shape, model_hash=model_hash)31latents = latents.to(device=pipe.device, dtype=torch.float16)32 33# generation without watermarking34image = pipe(prompt="an astronaut", latents=latents).images[0]35 36is_watermarked = trk.detect(image, pipe, model_hash)37print(f'is_watermarked: {is_watermarked}')38