Skip to content

Commit 0ca4b17

Browse files
author
Huy Vuong
committed
fix(mflux): adjust __call__ method to handle negative_prompt correctly
- Updated the __call__ method in BaseImageModel to extract the negative_prompt from kwargs, ensuring it is handled properly during image generation. - Modified the __call__ method in ImageGenerationModel to maintain compatibility with the updated parameter handling.
1 parent d6d65fa commit 0ca4b17

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

app/models/mflux.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -202,10 +202,12 @@ def _generate_image(self, prompt: str, negative_prompt: Optional[str], seed: int
202202
except Exception as e:
203203
raise ModelGenerationError(f"{self.__class__.__name__} generation failed: {e}") from e
204204

205-
def __call__(self, prompt: str, negative_prompt: Optional[str] = None, seed: int = 42, **kwargs) -> Image.Image:
205+
def __call__(self, prompt: str, seed: int = 42, **kwargs) -> Image.Image:
206206
"""Generate an image from a text prompt."""
207207
if not self._is_loaded:
208208
raise ModelLoadError("Model is not loaded. Cannot generate image.")
209+
210+
negative_prompt = kwargs.pop('negative_prompt', None)
209211

210212
# Validate inputs
211213
if not prompt or not prompt.strip():
@@ -428,7 +430,7 @@ def __init__(self, model_path: str, config_name: str, quantize: int = 8,
428430

429431
def __call__(self, prompt: str, seed: int = 42, **kwargs) -> Image.Image:
430432
"""Generate an image using the configured model."""
431-
return self.model_instance(prompt, seed, **kwargs)
433+
return self.model_instance(prompt, seed=seed, **kwargs)
432434

433435
@classmethod
434436
def get_available_configs(cls) -> list[str]:

0 commit comments

Comments
 (0)