bugfix: add spaces
Browse files- 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]
|