amazonaws-la commited on
Commit
0382e35
1 Parent(s): 10a0a96

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -2
app.py CHANGED
@@ -23,6 +23,7 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
23
  ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
24
  ENABLE_REFINER = os.getenv("ENABLE_REFINER", "1") == "1"
25
  ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
 
26
 
27
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
28
 
@@ -48,6 +49,7 @@ def generate(
48
  guidance_scale_refiner: float = 5.0,
49
  num_inference_steps_base: int = 25,
50
  num_inference_steps_refiner: int = 25,
 
51
  use_lora: bool = False,
52
  apply_refiner: bool = False,
53
  model = 'SG161222/Realistic_Vision_V6.0_B1_noVAE',
@@ -55,10 +57,16 @@ def generate(
55
  lora = 'amazonaws-la/juliette',
56
  ) -> PIL.Image.Image:
57
  if torch.cuda.is_available():
58
- vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
59
- pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
 
60
  pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
61
 
 
 
 
 
 
62
  if use_lora:
63
  pipe.load_lora_weights(lora)
64
  pipe.fuse_lora(lora_scale=0.7)
@@ -193,6 +201,7 @@ with gr.Blocks(css="style.css") as demo:
193
  step=32,
194
  value=1024,
195
  )
 
196
  use_lora = gr.Checkbox(label='Use Lora', value=False, visible=ENABLE_USE_LORA)
197
  apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
198
  with gr.Row():
@@ -255,6 +264,13 @@ with gr.Blocks(css="style.css") as demo:
255
  queue=False,
256
  api_name=False,
257
  )
 
 
 
 
 
 
 
258
  use_lora.change(
259
  fn=lambda x: gr.update(visible=x),
260
  inputs=use_lora,
@@ -300,6 +316,7 @@ with gr.Blocks(css="style.css") as demo:
300
  guidance_scale_refiner,
301
  num_inference_steps_base,
302
  num_inference_steps_refiner,
 
303
  use_lora,
304
  apply_refiner,
305
  model,
 
23
  ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
24
  ENABLE_REFINER = os.getenv("ENABLE_REFINER", "1") == "1"
25
  ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
26
+ ENABLE_USE_VAE = os.getenv("ENABLE_USE_VAE", "1") == "1"
27
 
28
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
29
 
 
49
  guidance_scale_refiner: float = 5.0,
50
  num_inference_steps_base: int = 25,
51
  num_inference_steps_refiner: int = 25,
52
+ use_vae: bool = False,
53
  use_lora: bool = False,
54
  apply_refiner: bool = False,
55
  model = 'SG161222/Realistic_Vision_V6.0_B1_noVAE',
 
57
  lora = 'amazonaws-la/juliette',
58
  ) -> PIL.Image.Image:
59
  if torch.cuda.is_available():
60
+
61
+ if not use_vae:
62
+ pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
63
  pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
64
 
65
+ if use_vae:
66
+ vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
67
+ pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
68
+ pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
69
+
70
  if use_lora:
71
  pipe.load_lora_weights(lora)
72
  pipe.fuse_lora(lora_scale=0.7)
 
201
  step=32,
202
  value=1024,
203
  )
204
+ use_vae = gr.Checkbox(label='Use VAE', value=False, visible=ENABLE_USE_VAE)
205
  use_lora = gr.Checkbox(label='Use Lora', value=False, visible=ENABLE_USE_LORA)
206
  apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
207
  with gr.Row():
 
264
  queue=False,
265
  api_name=False,
266
  )
267
+ use_vae.change(
268
+ fn=lambda x: gr.update(visible=x),
269
+ inputs=use_vae,
270
+ outputs=vaecall,
271
+ queue=False,
272
+ api_name=False,
273
+ )
274
  use_lora.change(
275
  fn=lambda x: gr.update(visible=x),
276
  inputs=use_lora,
 
316
  guidance_scale_refiner,
317
  num_inference_steps_base,
318
  num_inference_steps_refiner,
319
+ use_vae,
320
  use_lora,
321
  apply_refiner,
322
  model,