File size: 2,726 Bytes
b091cae
74073af
 
 
b091cae
6c40a22
b091cae
 
74073af
b091cae
74073af
 
 
 
 
 
b091cae
74073af
b091cae
 
 
74073af
b091cae
 
 
74073af
 
 
 
9e0f76e
74073af
b091cae
74073af
 
 
 
b091cae
74073af
 
b091cae
 
 
74073af
 
 
 
 
 
b091cae
 
 
 
 
74073af
 
 
 
 
 
 
 
 
 
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
import torch
from PIL.Image import Image
from diffusers import StableDiffusionXLPipeline
import torch.nn as nn
from pipelines.models import TextToImageRequest
from diffusers import DDIMScheduler, KDPM2DiscreteScheduler, DPMSolverSinglestepScheduler, 	HeunDiscreteScheduler
from torch import Generator
from loss import SchedulerWrapper
from onediffx import compile_pipe, load_pipe

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.1

  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")

    # quantizer = SDXLQuantization(pipeline.unet, 14)
    # quantizer.quantize_model()

    pipeline.scheduler = SchedulerWrapper(DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config))

    pipeline = compile_pipe(pipeline)
    load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--RobertML--cached-pipe-02/snapshots/58d70deae87034cce351b780b48841f9746d4ad7")

    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:
    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]