Update app.py
Browse files
app.py
CHANGED
@@ -1,6 +1,9 @@
|
|
1 |
# -*- coding: utf-8 -*-
|
2 |
"""
|
3 |
-
|
|
|
|
|
|
|
4 |
"""
|
5 |
|
6 |
import gradio as gr
|
@@ -102,10 +105,15 @@ if INITIAL_MODEL_ID:
|
|
102 |
)
|
103 |
|
104 |
# --- Apply Optimizations during initial load ---
|
105 |
-
# Apply attention slicing by default for memory efficiency on Spaces
|
106 |
-
|
107 |
-
|
108 |
-
|
|
|
|
|
|
|
|
|
|
|
109 |
|
110 |
pipeline = pipeline.to(initial_device_to_use) # Move to the initial device
|
111 |
|
@@ -155,8 +163,9 @@ def infer(
|
|
155 |
size, # From size_dropdown
|
156 |
seed, # From seed_input (now a Slider)
|
157 |
randomize_seed, # From randomize_seed_checkbox
|
158 |
-
enable_attention_slicing, #
|
159 |
-
|
|
|
160 |
):
|
161 |
"""Generates an image using the selected model and parameters on the chosen device."""
|
162 |
global current_pipeline, current_model_id, current_device_loaded, SCHEDULER_MAP, MAX_SEED
|
@@ -213,20 +222,20 @@ def infer(
|
|
213 |
if enable_attention_slicing and temp_device_to_use == "cuda": # Only apply on GPU
|
214 |
try:
|
215 |
pipeline.enable_attention_slicing()
|
216 |
-
print("Attention Slicing enabled.")
|
217 |
except Exception as e:
|
218 |
print(f"Warning: Failed to enable Attention Slicing: {e}")
|
219 |
gr.Warning(f"Failed to enable Attention Slicing. Error: {e}")
|
220 |
else:
|
221 |
try:
|
222 |
pipeline.disable_attention_slicing() # Ensure it's off if toggle is off or on CPU
|
223 |
-
print("Attention Slicing disabled.")
|
224 |
except Exception as e:
|
225 |
# May fail if it wasn't enabled, ignore
|
226 |
pass
|
227 |
|
228 |
|
229 |
-
pipeline = pipeline.to(temp_device_to_use) #
|
230 |
|
231 |
current_pipeline = pipeline
|
232 |
current_model_id = model_identifier
|
@@ -285,6 +294,7 @@ def infer(
|
|
285 |
|
286 |
# --- Apply Optimizations *before* generation if model was already loaded ---
|
287 |
# If the model didn't need reloading, we need to apply/remove slicing here
|
|
|
288 |
if str(current_pipeline.device) == "cuda": # Only attempt on GPU
|
289 |
if enable_attention_slicing:
|
290 |
try:
|
@@ -374,7 +384,7 @@ def infer(
|
|
374 |
print(f"Warning: Invalid seed input '{seed}'. Using random seed instead.")
|
375 |
gr.Warning(f"Invalid seed input '{seed}'. Using random seed instead.")
|
376 |
seed_int = random.randint(0, MAX_SEED) # Fallback to random if input is not int
|
377 |
-
randomize_seed = True
|
378 |
|
379 |
|
380 |
generator = None # Initialize generator as None
|
@@ -536,11 +546,11 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: # Added Soft theme from
|
|
536 |
)
|
537 |
# Combine seed input and randomize checkbox
|
538 |
with gr.Row():
|
539 |
-
# Removed precision=0 from Slider
|
540 |
seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, interactive=True) # Use 0 as default, interactive initially
|
541 |
randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True) # Simplified label
|
542 |
|
543 |
-
# ---
|
544 |
with gr.Row():
|
545 |
# Default to enabled if GPU is available, otherwise off
|
546 |
default_slicing = True if "GPU" in AVAILABLE_DEVICES else False
|
@@ -584,7 +594,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: # Added Soft theme from
|
|
584 |
size_dropdown,
|
585 |
seed_input,
|
586 |
randomize_seed_checkbox,
|
587 |
-
enable_attention_slicing_checkbox, #
|
588 |
],
|
589 |
outputs=[output_image, actual_seed_output], # Return image and the actual seed used
|
590 |
api_name="generate" # Optional: For API access
|
|
|
1 |
# -*- coding: utf-8 -*-
|
2 |
"""
|
3 |
+
Hugging Face Spaces Script for Basic Stable Diffusion 1.5 Gradio App
|
4 |
+
Adapted from user's local script and HF Spaces template.
|
5 |
+
Supports Hub models and CPU/GPU selection based on available hardware.
|
6 |
+
Includes Attention Slicing optimization toggle.
|
7 |
"""
|
8 |
|
9 |
import gradio as gr
|
|
|
105 |
)
|
106 |
|
107 |
# --- Apply Optimizations during initial load ---
|
108 |
+
# Apply attention slicing by default if GPU is available for memory efficiency on Spaces
|
109 |
+
if initial_device_to_use == "cuda":
|
110 |
+
try:
|
111 |
+
pipeline.enable_attention_slicing()
|
112 |
+
print("Attention Slicing enabled during initial load.")
|
113 |
+
except Exception as e:
|
114 |
+
print(f"Warning: Failed to enable Attention Slicing during initial load: {e}")
|
115 |
+
# Don't raise Gradio error here, just print warning
|
116 |
+
|
117 |
|
118 |
pipeline = pipeline.to(initial_device_to_use) # Move to the initial device
|
119 |
|
|
|
163 |
size, # From size_dropdown
|
164 |
seed, # From seed_input (now a Slider)
|
165 |
randomize_seed, # From randomize_seed_checkbox
|
166 |
+
enable_attention_slicing, # New input for the optimization toggle
|
167 |
+
# Removed track_tqdm=True which can conflict with callbacks
|
168 |
+
progress=gr.Progress(), # <-- Corrected Progress initialization
|
169 |
):
|
170 |
"""Generates an image using the selected model and parameters on the chosen device."""
|
171 |
global current_pipeline, current_model_id, current_device_loaded, SCHEDULER_MAP, MAX_SEED
|
|
|
222 |
if enable_attention_slicing and temp_device_to_use == "cuda": # Only apply on GPU
|
223 |
try:
|
224 |
pipeline.enable_attention_slicing()
|
225 |
+
print("Attention Slicing enabled during model load.")
|
226 |
except Exception as e:
|
227 |
print(f"Warning: Failed to enable Attention Slicing: {e}")
|
228 |
gr.Warning(f"Failed to enable Attention Slicing. Error: {e}")
|
229 |
else:
|
230 |
try:
|
231 |
pipeline.disable_attention_slicing() # Ensure it's off if toggle is off or on CPU
|
232 |
+
# print("Attention Slicing disabled during model load.") # Avoid noise
|
233 |
except Exception as e:
|
234 |
# May fail if it wasn't enabled, ignore
|
235 |
pass
|
236 |
|
237 |
|
238 |
+
pipeline = pipeline.to(temp_device_to_use) # Move to the determined device
|
239 |
|
240 |
current_pipeline = pipeline
|
241 |
current_model_id = model_identifier
|
|
|
294 |
|
295 |
# --- Apply Optimizations *before* generation if model was already loaded ---
|
296 |
# If the model didn't need reloading, we need to apply/remove slicing here
|
297 |
+
# Check the pipeline's actual device
|
298 |
if str(current_pipeline.device) == "cuda": # Only attempt on GPU
|
299 |
if enable_attention_slicing:
|
300 |
try:
|
|
|
384 |
print(f"Warning: Invalid seed input '{seed}'. Using random seed instead.")
|
385 |
gr.Warning(f"Invalid seed input '{seed}'. Using random seed instead.")
|
386 |
seed_int = random.randint(0, MAX_SEED) # Fallback to random if input is not int
|
387 |
+
# No need to set randomize_seed = True here, just use the new random seed_int
|
388 |
|
389 |
|
390 |
generator = None # Initialize generator as None
|
|
|
546 |
)
|
547 |
# Combine seed input and randomize checkbox
|
548 |
with gr.Row():
|
549 |
+
# Removed precision=0 from Slider - FIX FOR TYPEERROR
|
550 |
seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, interactive=True) # Use 0 as default, interactive initially
|
551 |
randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True) # Simplified label
|
552 |
|
553 |
+
# --- Memory Optimization Toggle ---
|
554 |
with gr.Row():
|
555 |
# Default to enabled if GPU is available, otherwise off
|
556 |
default_slicing = True if "GPU" in AVAILABLE_DEVICES else False
|
|
|
594 |
size_dropdown,
|
595 |
seed_input,
|
596 |
randomize_seed_checkbox,
|
597 |
+
enable_attention_slicing_checkbox, # Pass the new checkbox value
|
598 |
],
|
599 |
outputs=[output_image, actual_seed_output], # Return image and the actual seed used
|
600 |
api_name="generate" # Optional: For API access
|