LPX55 commited on
Commit
85f921d
·
1 Parent(s): 30f6115

♻️ refactor(app): improve pipeline loading mechanism

Browse files

- move pipeline loading into a separate function `load_pipeline` for lazy loading
- use a global `pipe` variable and update it dynamically based on the selected model
- remove redundant pipeline loading code and use the `load_pipeline` function instead
- improve code organization and readability by separating concerns into different functions

Files changed (1) hide show
  1. app.py +14 -25
app.py CHANGED
@@ -39,25 +39,17 @@ vae = AutoencoderKL.from_pretrained(
39
  "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
40
  ).to("cuda")
41
 
42
- pipe = StableDiffusionXLFillPipeline.from_pretrained(
43
- "SG161222/RealVisXL_V5.0_Lightning",
44
- torch_dtype=torch.float16,
45
- vae=vae,
46
- controlnet=model,
47
- variant="fp16",
48
- )
49
-
50
- pipe = StableDiffusionXLFillPipeline.from_pretrained(
51
- "GraydientPlatformAPI/lustify-lightning",
52
- torch_dtype=torch.float16,
53
- vae=vae,
54
- controlnet=model,
55
- )
56
-
57
- pipe.to("cuda")
58
-
59
- pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
60
-
61
 
62
  def can_expand(source_width, source_height, target_width, target_height, alignment):
63
  """Checks if the image can be expanded based on the alignment."""
@@ -187,12 +179,8 @@ def preview_image_and_mask(image, width, height, overlap_percentage, resize_opti
187
  def inpaint(prompt, image, model_name, paste_back):
188
  global pipe
189
  if pipe.config.model_name != MODELS[model_name]:
190
- pipe = StableDiffusionXLFillPipeline.from_pretrained(
191
- MODELS[model_name],
192
- torch_dtype=torch.float16,
193
- vae=vae,
194
- controlnet=model,
195
- ).to("cuda")
196
 
197
  mask = Image.fromarray(image["mask"]).convert("L")
198
  image = Image.fromarray(image["image"])
@@ -206,6 +194,7 @@ def inpaint(prompt, image, model_name, paste_back):
206
 
207
  @spaces.GPU(duration=24)
208
  def outpaint(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
 
209
  background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
210
 
211
  if not can_expand(background.width, background.height, width, height, alignment):
 
39
  "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
40
  ).to("cuda")
41
 
42
+ # Move pipeline loading into a function to enable lazy loading
43
+ def load_pipeline(model_name):
44
+ pipe = StableDiffusionXLFillPipeline.from_pretrained(
45
+ MODELS[model_name],
46
+ torch_dtype=torch.float16,
47
+ vae=vae,
48
+ controlnet=model,
49
+ )
50
+ pipe.to("cuda")
51
+ pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
52
+ return pipe
 
 
 
 
 
 
 
 
53
 
54
  def can_expand(source_width, source_height, target_width, target_height, alignment):
55
  """Checks if the image can be expanded based on the alignment."""
 
179
  def inpaint(prompt, image, model_name, paste_back):
180
  global pipe
181
  if pipe.config.model_name != MODELS[model_name]:
182
+ # Lazily load the pipeline for the selected model
183
+ pipe = load_pipeline(model_name)
 
 
 
 
184
 
185
  mask = Image.fromarray(image["mask"]).convert("L")
186
  image = Image.fromarray(image["image"])
 
194
 
195
  @spaces.GPU(duration=24)
196
  def outpaint(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
197
+ # Use the currently loaded pipeline
198
  background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
199
 
200
  if not can_expand(background.width, background.height, width, height, alignment):