bugfix: fix params
Browse files- inference.py +3 -4
inference.py
CHANGED
@@ -37,7 +37,6 @@ class DiffusionInference:
|
|
37 |
# Create parameters dictionary for all keyword arguments
|
38 |
params = {
|
39 |
"prompt": prompt,
|
40 |
-
"model": model
|
41 |
}
|
42 |
|
43 |
# Handle seed parameter
|
@@ -67,7 +66,7 @@ class DiffusionInference:
|
|
67 |
|
68 |
try:
|
69 |
# Call the API with all parameters as kwargs
|
70 |
-
image = self.run_text_to_image_pipeline(**params)
|
71 |
return image
|
72 |
except Exception as e:
|
73 |
print(f"Error generating image: {e}")
|
@@ -164,7 +163,7 @@ class DiffusionInference:
|
|
164 |
print(f"Warning: Could not delete temporary file {temp_file}: {e}")
|
165 |
|
166 |
@spaces.GPU
|
167 |
-
def run_text_to_image_pipeline(self, **kwargs):
|
168 |
-
pipeline = AutoPipelineForText2Image.from_pretrained(model=
|
169 |
image = pipeline(**kwargs).images[0]
|
170 |
return image
|
|
|
37 |
# Create parameters dictionary for all keyword arguments
|
38 |
params = {
|
39 |
"prompt": prompt,
|
|
|
40 |
}
|
41 |
|
42 |
# Handle seed parameter
|
|
|
66 |
|
67 |
try:
|
68 |
# Call the API with all parameters as kwargs
|
69 |
+
image = self.run_text_to_image_pipeline(model, **params)
|
70 |
return image
|
71 |
except Exception as e:
|
72 |
print(f"Error generating image: {e}")
|
|
|
163 |
print(f"Warning: Could not delete temporary file {temp_file}: {e}")
|
164 |
|
165 |
@spaces.GPU
|
166 |
+
def run_text_to_image_pipeline(self, model_name, **kwargs):
|
167 |
+
pipeline = AutoPipelineForText2Image.from_pretrained(model=model_name, torch_dtype=torch.float16).to("cuda")
|
168 |
image = pipeline(**kwargs).images[0]
|
169 |
return image
|