Split the models into 2

src/pipeline.py

@@ -1,10 +1,11 @@
 import torch
 from PIL.Image import Image
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionXLPipeline, AutoPipelineForImage2Image
 from pipelines.models import TextToImageRequest
 from torch import Generator
 from DeepCache import DeepCacheSDHelper
 
+'''
 def callback_dynamic_cfg(pipeline, step_index, timestep, callback_kwargs):
   if step_index == int(pipeline.num_timesteps * 0.5):
     callback_kwargs['prompt_embeds'] = callback_kwargs['prompt_embeds'].chunk(2)[-1]
@@ -13,10 +14,13 @@ def callback_dynamic_cfg(pipeline, step_index, timestep, callback_kwargs):
     pipeline._guidance_scale = 0.0
 
   return callback_kwargs
-
+'''
 
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-def load_pipeline() -> StableDiffusionXLPipeline:
+
+def load_pipeline() -> dict: #StableDiffusionXLPipeline, AutoPipelineForImage2Image:
+    
+    pipeline_dict = {}
     pipeline = StableDiffusionXLPipeline.from_pretrained(
         "./models/newdream-sdxl-20",
         torch_dtype=torch.float16,
@@ -25,22 +29,63 @@ def load_pipeline() -> StableDiffusionXLPipeline:
         variant='fp16',        
     ).to("cuda")
 
+    refiner = AutoPipelineForImage2Image.from_pretrained(
+        'stabilityai/stable-diffusion-xl-refiner-1.0',
+        use_safetensors=True,
+        torch_dtype=torch.float16,
+        variant='fp16',
+        ).to('cuda')
+
+    
     helper = DeepCacheSDHelper(pipe=pipeline)
     helper.set_params(cache_interval=3, cache_branch_id=0)
     helper.enable()    
 
+    refiner_helper = DeepCacheSDHelper(pipe=refiner)
+    refiner_helper.set_params(cache_interval=3, cache_branch_id=0)
+    refiner_helper.enable()    
+    
+
     for _ in range(5):
         pipeline(prompt="")
 
-    return pipeline
+    pipeline_dict = {
+        'base_pipeline': pipeline,
+        'refiner': refiner
+    }
+    return pipeline_dict #base_pipeline, refiner
 
 
-def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
+def infer(request: TextToImageRequest, pipeline_dict: dict) -> Image: #pipeline: StableDiffusionXLPipeline, refiner: AutoPipelineForImage2Image) -> Image:
     if request.seed is None:
         generator = None
     else:
-        generator = Generator(pipeline.device).manual_seed(request.seed)
+        generator = Generator(pipeline_dict['base_pipeline'].device).manual_seed(request.seed)
+
+
+    image = pipeline_dict['base_pipeline'](
+        prompt=request.prompt,
+        negative_prompt=request.negative_prompt,
+        width=request.width,
+        height=request.height,
+        generator=generator,
+        num_inference_steps=27,
+        denoising_end=0.8,
+        output_type='latent',
+    ).images
+
+    return pipeline_dict['refiner'](
+        prompt=request.prompt,
+        negative_prompt=request.negative_prompt,
+        width=request.width,
+        height=request.height,
+        generator=generator,
+        num_inference_steps=27,
+    denoising_start=0.8,
+    image=image,
+  ).images[0]
 
+    '''
     return pipeline(
         prompt=request.prompt,
         negative_prompt=request.negative_prompt,
@@ -48,6 +93,7 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
         height=request.height,
         generator=generator,
         num_inference_steps=27,
-        callback_on_step_end=callback_dynamic_cfg,
-        callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],        
+        #callback_on_step_end=callback_dynamic_cfg,
+        #callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],        
     ).images[0]
+    '''