File size: 4,394 Bytes
80b58c8 a44493b 80b58c8 a44493b 80b58c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """
Byte Dream - Command Line Inference Tool
Generate images from text prompts using the command line
"""
import argparse
from pathlib import Path
import torch
def main():
parser = argparse.ArgumentParser(
description="Byte Dream - AI Image Generation",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Basic usage
python infer.py --prompt "A beautiful sunset over mountains"
# With custom parameters
python infer.py --prompt "Cyberpunk city" --negative "blurry" --steps 75 --guidance 8.0
# Specify output and size
python infer.py --prompt "Fantasy landscape" --output fantasy.png --width 768 --height 768
# With seed for reproducibility
python infer.py --prompt "Dragon" --seed 42 --output dragon.png
"""
)
parser.add_argument(
"--prompt", "-p",
type=str,
required=True,
help="Text prompt describing the desired image"
)
parser.add_argument(
"--negative", "-n",
type=str,
default="",
help="Negative prompt - things to avoid in the image"
)
parser.add_argument(
"--output", "-o",
type=str,
default="output.png",
help="Output image filename (default: output.png)"
)
parser.add_argument(
"--width", "-W",
type=int,
default=512,
help="Image width in pixels (default: 512)"
)
parser.add_argument(
"--height", "-H",
type=int,
default=512,
help="Image height in pixels (default: 512)"
)
parser.add_argument(
"--steps", "-s",
type=int,
default=50,
help="Number of inference steps (default: 50)"
)
parser.add_argument(
"--guidance", "-g",
type=float,
default=7.5,
help="Guidance scale - how closely to follow prompt (default: 7.5)"
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for reproducibility (default: random)"
)
parser.add_argument(
"--model", "-m",
type=str,
default=None,
help="Path to model directory or Hugging Face repo ID (default: uses config)"
)
parser.add_argument(
"--hf_repo",
type=str,
default=None,
help="Hugging Face repository ID to load model from (e.g., username/repo)"
)
parser.add_argument(
"--config", "-c",
type=str,
default="config.yaml",
help="Path to config file (default: config.yaml)"
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Device to run on: cpu or cuda (default: cpu)"
)
args = parser.parse_args()
# Import generator
from bytedream.generator import ByteDreamGenerator
# Initialize generator
print("="*60)
print("Byte Dream - AI Image Generator")
print("="*60)
# Determine if loading from HF or local
if args.hf_repo:
print(f"Loading model from Hugging Face: {args.hf_repo}")
generator = ByteDreamGenerator(
hf_repo_id=args.hf_repo,
config_path=args.config,
device=args.device,
)
else:
generator = ByteDreamGenerator(
model_path=args.model,
config_path=args.config,
device=args.device,
)
# Print model info
info = generator.get_model_info()
print(f"\nModel: {info['name']} v{info['version']}")
print(f"Device: {info['device']}")
print(f"Parameters: {info['unet_parameters']}")
print("="*60)
# Generate image
image = generator.generate(
prompt=args.prompt,
negative_prompt=args.negative if args.negative else None,
width=args.width,
height=args.height,
num_inference_steps=args.steps,
guidance_scale=args.guidance,
seed=args.seed,
)
# Save image
output_path = Path(args.output)
image.save(output_path)
print(f"\n✓ Image saved to: {output_path.absolute()}")
print("="*60)
if __name__ == "__main__":
main()
|