bugfix: fix params
Browse files- inference.py +1 -1
inference.py
CHANGED
@@ -164,6 +164,6 @@ class DiffusionInference:
|
|
164 |
|
165 |
@spaces.GPU
|
166 |
def run_text_to_image_pipeline(self, model_name, **kwargs):
|
167 |
-
pipeline = AutoPipelineForText2Image.from_pretrained(
|
168 |
image = pipeline(**kwargs).images[0]
|
169 |
return image
|
|
|
164 |
|
165 |
@spaces.GPU
|
166 |
def run_text_to_image_pipeline(self, model_name, **kwargs):
|
167 |
+
pipeline = AutoPipelineForText2Image.from_pretrained(model_name, torch_dtype=torch.float16).to("cuda")
|
168 |
image = pipeline(**kwargs).images[0]
|
169 |
return image
|