Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 =
|
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=
|
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(
|