File size: 2,482 Bytes
335ad0f
d6248de
 
335ad0f
d6248de
335ad0f
 
d6248de
335ad0f
d6248de
356fdec
d6248de
 
 
356fdec
335ad0f
d6248de
335ad0f
 
 
d6248de
335ad0f
d6248de
 
 
335ad0f
 
d6248de
 
335ad0f
d6248de
 
335ad0f
 
 
d6248de
 
 
 
335ad0f
d6248de
335ad0f
 
 
 
 
d6248de
 
 
 
 
356fdec
d6248de
 
 
 
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
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 onediffx import compile_pipe, save_pipe, load_pipe

def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
  if step_index == int(pipe.num_timesteps * 0.68):
    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,)
    pipeline.to("cuda")
    pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))    
    pipeline = compile_pipe(pipeline)
    load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--slobers--make_me_happy/cucumber")
    for _ in range(1):
        deepcache_output = pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", output_type="pil", cache_interval=1, cache_layer_id=1, cache_block_id=0)
    pipeline.scheduler.prepare_loss()
    for _ in range(2):
        pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", output_type="pil", num_inference_steps=20)
    return pipeline

def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
    if request.seed is None:
        generator = None
    else:
        generator = Generator(pipeline.device).manual_seed(request.seed)

    return pipeline(
        prompt=request.prompt,
        negative_prompt=request.negative_prompt,
        width=request.width,
        height=request.height,
        generator=generator,
        num_inference_steps=13,
        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]