ovedrive commited on
Commit
6d82cf3
·
1 Parent(s): aff01bb

bugfix: fix params

Browse files
Files changed (1) hide show
  1. 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=model, torch_dtype=torch.float16).to("cuda")
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