Raxephion commited on
Commit
d7ff0be
·
verified ·
1 Parent(s): ce1167c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -14
app.py CHANGED
@@ -1,6 +1,9 @@
1
  # -*- coding: utf-8 -*-
2
  """
3
- Author: @Raxephion 2025
 
 
 
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
- # Can be turned off via UI toggle later, but good default for VRAM
107
- # We'll add the UI toggle later, for now, just enable it here
108
- # pipeline.enable_attention_slicing() # Enable by default on initial load
 
 
 
 
 
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, # <-- New input for the optimization toggle
159
- progress=gr.Progress(track_tqdm=True), # Added progress argument from template
 
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) # Use the determined device
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 # Mark as randomized for reporting
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
- # --- New: Memory Optimization Toggle ---
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, # <-- Pass the new checkbox value
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