ovedrive commited on
Commit
527eee1
·
1 Parent(s): 205fcdc

bugfix: add spaces

Browse files
Files changed (1) hide show
  1. inference.py +2 -0
inference.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from huggingface_hub import InferenceClient
2
  from PIL import Image
3
  import io
@@ -162,6 +163,7 @@ class DiffusionInference:
162
  except Exception as e:
163
  print(f"Warning: Could not delete temporary file {temp_file}: {e}")
164
 
 
165
  def run_text_to_image_pipeline(self,model_name, **kwargs):
166
  pipeline = AutoPipelineForText2Image.from_pretrained(model=model_name, torch_dtype=torch.float16).to("cuda")
167
  image = pipeline(**kwargs).images[0]
 
1
+ import spaces
2
  from huggingface_hub import InferenceClient
3
  from PIL import Image
4
  import io
 
163
  except Exception as e:
164
  print(f"Warning: Could not delete temporary file {temp_file}: {e}")
165
 
166
+ @spaces.GPU
167
  def run_text_to_image_pipeline(self,model_name, **kwargs):
168
  pipeline = AutoPipelineForText2Image.from_pretrained(model=model_name, torch_dtype=torch.float16).to("cuda")
169
  image = pipeline(**kwargs).images[0]