Views
No views yet
1Pythonimport torch
2from diffusers import StableDiffusionPipeline, DDIMScheduler
3from scheduler_ppo import PPOScheduler # Provided in this repo
4from huggingface_hub import hf_hub_download
5
6# Download the trained factor_net checkpoint
7factor_net_path = hf_hub_download(
8 repo_id="wangfuyun/consolver",
9 filename="model.ckpt"
10)
11
12model_id = "runwayml/stable-diffusion-v1-5"
13prompt = "an astronaut riding a horse on the moon, highly detailed, 8k"
14num_inference_steps = 8
15guidance_scale = 3.0
16seed = 43
17height = width = 512
18
19def load_pipeline(scheduler_type="ddim"):
20 if scheduler_type == "ppo":
21 scheduler = PPOScheduler(
22 beta_end=0.012,
23 beta_schedule="scaled_linear",
24 beta_start=0.00085,
25 num_train_timesteps=1000,
26 steps_offset=1,
27 timestep_spacing="trailing",
28 order_dim=4,
29 scaler_dim=0,
30 use_conv=False,
31 factor_net_kwargs=dict(embedding_dim=64, hidden_dim=256, num_actions=11),
32 )
33 else:
34 scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler", timestep_spacing="trailing")
35
36 pipe = StableDiffusionPipeline.from_pretrained(
37 model_id,
38 scheduler=scheduler,
39 safety_checker=None,
40 # torch_dtype=torch.float16, # Uncomment for GPU memory savings
41 ).to("cuda")
42
43 if scheduler_type == "ppo" and factor_net_path:
44 weight = torch.load(factor_net_path, map_location="cpu")
45 pipe.scheduler.factor_net.load_state_dict(weight)
46 pipe.scheduler.factor_net.to("cuda")
47
48 return pipe
49
50generator = torch.Generator("cuda").manual_seed(seed)
51
52# DDIM baseline (8 steps)
53pipe_ddim = load_pipeline("ddim")
54image_ddim = pipe_ddim(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale,
55 generator=generator, height=height, width=width).images[0]
56image_ddim.save("ddim_result.jpg")
57
58# ConSolver (8 steps)
59pipe_consolver = load_pipeline("ppo")
60image_consolver = pipe_consolver(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale,
61 generator=generator, height=height, width=width).images[0]
62image_consolver.save("consolver_result.jpg")|
|
|
| DDIM | ConsistencySolver |