File size: 3,079 Bytes
a5399c0 ce15133 a5399c0 14ea712 a5399c0 5590c65 a5399c0 29bbbb9 fe21164 1588824 61b487f fe21164 14ea712 1588824 a5399c0 335ee70 73f600a a5399c0 1588824 a5399c0 14ea712 382952a 14ea712 4976401 a5399c0 2bb06b3 73f600a 5590c65 73f600a a5399c0 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import torch
from PIL.Image import Image
from diffusers import StableDiffusionXLPipeline
from pipelines.models import TextToImageRequest
from diffusers import DDIMScheduler
from torch import Generator
from loss import SchedulerWrapper
from utils import register_parallel_pipeline_orig, register_faster_orig_forward, register_time
from onediffx import compile_pipe, save_pipe, load_pipe
import numpy as np
def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
if step_index == int(pipe.num_timesteps * 0.78):
callback_kwargs['prompt_embeds'] = callback_kwargs['prompt_embeds'].chunk(2)[-1]
callback_kwargs['add_text_embeds'] = callback_kwargs['add_text_embeds'].chunk(2)[-1]
callback_kwargs['add_time_ids'] = callback_kwargs['add_time_ids'].chunk(2)[-1]
pipe._guidance_scale = 0.0
return callback_kwargs
def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
if not pipeline:
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stablediffusionapi/newdream-sdxl-20",
torch_dtype=torch.float16,
).to("cuda")
pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))
register_parallel_pipeline_orig(pipeline)
register_faster_orig_forward(pipeline.unet)
register_time(pipeline.unet, 0)
pipeline = compile_pipe(pipeline)
# load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--RobertML--cached-pipe-02/snapshots/58d70deae87034cce351b780b48841f9746d4ad7")
np.save('count_array.npy', [0])
for _ in range(1):
deepcache_output = pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
pipeline.scheduler.prepare_loss()
# for _ in range(2):
# pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
return pipeline
def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
global image_count
if request.seed is None:
generator = None
else:
generator = Generator(pipeline.device).manual_seed(request.seed)
step = np.load('count_array.npy')[0]
register_time(pipeline.unet, step)
print("step is : ", step)
step+=1; np.save('count_array.npy', [step])
# register_time(pipeline.unet, image_count)
return pipeline(
prompt=request.prompt,
negative_prompt=request.negative_prompt,
width=request.width,
height=request.height,
generator=generator,
num_inference_steps=20,
# cache_interval=1,
# cache_layer_id=1,
# cache_block_id=0,
eta=1.0,
guidance_scale = 5.0,
guidance_rescale = 0.0,
# callback_on_step_end=callback_dynamic_cfg,
# callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],
).images[0]
|