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]
|