kaupane commited on
Commit
61df749
·
verified ·
1 Parent(s): eb3157d

Update app.py

Browse files

Add spaces to utilize zeroGPU.

Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import torch
 
2
  import gradio as gr
3
  import numpy as np
4
  import os
@@ -65,7 +66,8 @@ class DiffusionSampler:
65
  if self.vae is None:
66
  self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema").to(self.device)
67
  self.vae.eval()
68
-
 
69
  def generate_images(self, model, num_samples, genre, style, seed, progress=gr.Progress()):
70
  """Generate images with the DiT model"""
71
  global global_progress
@@ -159,6 +161,7 @@ def generate_random_seed():
159
  """Generate a random seed between 0 and 2^32 - 1"""
160
  return random.randint(0, 2**32 - 1)
161
 
 
162
  def generate_samples(num_samples, dit_size, genre_name, style_name, seed, progress=gr.Progress()):
163
  """Main function for Gradio interface"""
164
  if num_samples < 1 or num_samples > 16:
 
1
  import torch
2
+ import spaces
3
  import gradio as gr
4
  import numpy as np
5
  import os
 
66
  if self.vae is None:
67
  self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema").to(self.device)
68
  self.vae.eval()
69
+
70
+ @spaces.GPU
71
  def generate_images(self, model, num_samples, genre, style, seed, progress=gr.Progress()):
72
  """Generate images with the DiT model"""
73
  global global_progress
 
161
  """Generate a random seed between 0 and 2^32 - 1"""
162
  return random.randint(0, 2**32 - 1)
163
 
164
+ @spaces.GPU
165
  def generate_samples(num_samples, dit_size, genre_name, style_name, seed, progress=gr.Progress()):
166
  """Main function for Gradio interface"""
167
  if num_samples < 1 or num_samples > 16: