|
#!/usr/bin/env python3 |
|
import argparse |
|
import os |
|
from pathlib import Path |
|
import subprocess |
|
|
|
def generate_image(prompt: str, output_dir: str = "outputs", num_steps: int = 2, |
|
height: int = 1024, width: int = 1024, seed: int = 42) -> str: |
|
"""Generate an image from a text prompt using MFLUX""" |
|
# Ensure output directory exists |
|
Path(output_dir).mkdir(parents=True, exist_ok=True) |
|
|
|
# Generate a filename from the prompt |
|
safe_prompt = "".join(x for x in prompt[:30] if x.isalnum() or x in (' ', '-', '_')).strip() |
|
output_path = os.path.join(output_dir, f"{safe_prompt}.png") |
|
|
|
print(f"Generating image for prompt: {prompt}") |
|
print(f"Using {num_steps} inference steps") |
|
print(f"Output will be saved to: {output_path}") |
|
|
|
try: |
|
# Use mflux-generate command |
|
cmd = [ |
|
"mflux-generate", |
|
"--model", "schnell", # Using the schnell model as recommended |
|
"--steps", str(num_steps), |
|
"--height", str(height), |
|
"--width", str(width), |
|
"--seed", str(seed), |
|
"--output", output_path, |
|
"--prompt", prompt |
|
] |
|
|
|
# Run the command |
|
subprocess.run(cmd, check=True) |
|
|
|
if os.path.exists(output_path): |
|
print(f"Image generated successfully!") |
|
return output_path |
|
else: |
|
raise Exception("Output image not found") |
|
|
|
except Exception as e: |
|
print(f"Error generating image: {str(e)}") |
|
return None |
|
|
|
def main(): |
|
parser = argparse.ArgumentParser(description="Generate images using MFLUX") |
|
parser.add_argument("prompt", type=str, help="Text prompt for image generation") |
|
parser.add_argument("--steps", type=int, default=2, help="Number of inference steps") |
|
parser.add_argument("--output", type=str, default="outputs", help="Output directory") |
|
parser.add_argument("--height", type=int, default=1024, help="Image height") |
|
parser.add_argument("--width", type=int, default=1024, help="Image width") |
|
parser.add_argument("--seed", type=int, default=42, help="Random seed") |
|
|
|
args = parser.parse_args() |
|
|
|
output_path = generate_image( |
|
args.prompt, |
|
args.output, |
|
args.steps, |
|
args.height, |
|
args.width, |
|
args.seed |
|
) |
|
|
|
if output_path: |
|
print(f"\nImage saved to: {output_path}") |
|
else: |
|
print("\nFailed to generate image") |
|
|
|
if __name__ == "__main__": |
|
main() |