Wz-2024/carton

Usage

需要下载google的模型来构造一个Scheduler

import torch
from diffusers import DDIMPipeline,DDIMScheduler
import torch
import torchvision
from diffusers import DDIMScheduler 
from matplotlib import pyplot as plt
from tqdm.auto import tqdm
device='cuda:1'
x=torch.randn(8,3,128,128).to(device)
path='your_model_path'
# path='/data_disk/dyy/python_projects/diffusers/0.My_model/butterfly_generate'
image_pipe = DDIMPipeline.from_pretrained(path).to(device)
#这里要填一个DDPM或者DDIMbased模型,,高手可以自己构造一个timestep
scheduler=DDIMScheduler.from_pretrained('/data_disk/dyy/models/google-ddpm')
scheduler.set_timesteps(30)

for idx , t  in tqdm(enumerate(scheduler.timesteps)):
    model_input=scheduler.scale_model_input(x,t)
    with torch.no_grad():
        noise_pred=image_pipe.unet(model_input,t)['sample']
    x=scheduler.step(noise_pred,t,x).prev_sample.to(device)
grid=torchvision.utils.make_grid(x,nrow=4)
plt.imshow(grid.permute(1,2,0).cpu().clip(-1,1)*0.5+0.5)