KingNish commited on
Commit
8c0ff20
·
verified ·
1 Parent(s): 2aa1224

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -11
app.py CHANGED
@@ -57,21 +57,12 @@ try:
57
  logging.info("LoRA fused and unloaded.")
58
 
59
  # --- Compilation (Major Speed Optimization) ---
60
- # Note: Compilation takes time on the first run.
61
- # logging.info("Compiling UNet (this may take a moment)...")
62
- # pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) # Use reduce-overhead for dynamic shapes
63
  # logging.info("Compiling VAE Decoder...")
64
  # pipe.vae.decoder = torch.compile(pipe.vae.decoder, mode="reduce-overhead", fullgraph=True)
65
  # logging.info("Compiling VAE Encoder...")
66
  # pipe.vae.encoder = torch.compile(pipe.vae.encoder, mode="reduce-overhead", fullgraph=True)
67
  # logging.info("Model compilation finished.")
68
 
69
- # --- Optional: Warm-up Run ---
70
- # logging.info("Performing warm-up run...")
71
- # with torch.inference_mode():
72
- # _ = pipe(prompt="warmup", num_inference_steps=1, generator=torch.Generator(device=device).manual_seed(0), output_type="pil", return_dict=False)[0]
73
- # logging.info("Warm-up complete.")
74
-
75
  # Clear cache after setup
76
  if torch.cuda.is_available():
77
  torch.cuda.empty_cache()
@@ -87,7 +78,7 @@ except Exception as e:
87
 
88
  # --- Inference Function ---
89
  @spaces.GPU(duration=30) # Slightly increased duration buffer
90
- def generate_image(prompt: str, seed: int = 42, width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, randomize_seed: bool = False, num_inference_steps: int = 4, is_enhance: bool = False):
91
  """Generates an image using the FLUX pipeline with error handling."""
92
 
93
  if pipe is None:
@@ -224,7 +215,7 @@ with gr.Blocks() as demo:
224
  fn=generate_image,
225
  inputs=[prompt],
226
  outputs=[result, seed, latency],
227
- cache_examples="lazy"
228
  )
229
 
230
  enhanceBtn.click(
 
57
  logging.info("LoRA fused and unloaded.")
58
 
59
  # --- Compilation (Major Speed Optimization) ---
 
 
 
60
  # logging.info("Compiling VAE Decoder...")
61
  # pipe.vae.decoder = torch.compile(pipe.vae.decoder, mode="reduce-overhead", fullgraph=True)
62
  # logging.info("Compiling VAE Encoder...")
63
  # pipe.vae.encoder = torch.compile(pipe.vae.encoder, mode="reduce-overhead", fullgraph=True)
64
  # logging.info("Model compilation finished.")
65
 
 
 
 
 
 
 
66
  # Clear cache after setup
67
  if torch.cuda.is_available():
68
  torch.cuda.empty_cache()
 
78
 
79
  # --- Inference Function ---
80
  @spaces.GPU(duration=30) # Slightly increased duration buffer
81
+ def generate_image(prompt: str, seed: int = 42, width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, randomize_seed: bool = False, num_inference_steps: int = DEFAULT_INFERENCE_STEPS, is_enhance: bool = False):
82
  """Generates an image using the FLUX pipeline with error handling."""
83
 
84
  if pipe is None:
 
215
  fn=generate_image,
216
  inputs=[prompt],
217
  outputs=[result, seed, latency],
218
+ cache_examples=True
219
  )
220
 
221
  enhanceBtn.click(