diff --git a/.gitignore b/.gitignore index 550e762..4c689e1 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,4 @@ wheels/ .coverage htmlcov/ coverage.xml +models.db diff --git a/tensors/cli.py b/tensors/cli.py index ce03f8f..9ed37a7 100644 --- a/tensors/cli.py +++ b/tensors/cli.py @@ -1218,69 +1218,62 @@ def comfy_generate( # noqa: PLR0915 truncated = enhanced_negative[:80] + "..." if len(enhanced_negative) > 80 else enhanced_negative # noqa: PLR2004 console.print(f"[dim]Enhanced negative: {truncated}[/dim]") - for i in range(count): - current_seed = base_seed + i if seed >= 0 else -1 # Increment seed or use random each time + # Use native ComfyUI batching - single workflow generates all images + result = generate_image( + prompt=enhanced_prompt, + url=url, + negative_prompt=enhanced_negative, + model=model, + width=width, + height=height, + steps=steps, + cfg=cfg, + seed=base_seed, + sampler=sampler, + scheduler=scheduler, + console=console if not json_output else None, + lora_name=lora, + lora_strength=lora_strength, + batch_size=count, + ) - if count > 1 and not json_output: - console.print(f"\n[cyan]Generating image {i + 1}/{count}...[/cyan]") + if not result: + if json_output: + all_results.append({"success": False, "index": 0, "errors": {"generation": "Failed to generate"}}) + else: + console.print("[red]Generation failed[/red]") + elif not result.success: + if json_output: + all_results.append({"success": False, "index": 0, "errors": result.node_errors}) + else: + console.print("[red]Generation failed[/red]") + for node_id, errors in result.node_errors.items(): + console.print(f" [yellow]Node {node_id}:[/yellow] {errors}") + else: + # Save all output images + for i, img_path in enumerate(result.images): + saved_path: Path | None = None + if output: + img_data = get_image(str(img_path), url=url) + if img_data: + save_path = output if count == 1 else output.parent / f"{output.stem}_{i + 1:03d}{output.suffix}" + save_path.write_bytes(img_data) + saved_path = save_path + all_saved.append(save_path) + if not json_output: + console.print(f"[green]Saved:[/green] {save_path}") + elif not json_output: + console.print(f"[yellow]Could not download image: {img_path}[/yellow]") - result = generate_image( - prompt=enhanced_prompt, - url=url, - negative_prompt=enhanced_negative, - model=model, - width=width, - height=height, - steps=steps, - cfg=cfg, - seed=current_seed, - sampler=sampler, - scheduler=scheduler, - console=console if not json_output else None, - lora_name=lora, - lora_strength=lora_strength, - ) - - if not result: - if json_output: - all_results.append({"success": False, "index": i, "errors": {"generation": "Failed to generate"}}) - else: - console.print(f"[red]Generation {i + 1} failed[/red]") - continue - - if not result.success: - if json_output: - all_results.append({"success": False, "index": i, "errors": result.node_errors}) - else: - console.print(f"[red]Generation {i + 1} failed[/red]") - for node_id, errors in result.node_errors.items(): - console.print(f" [yellow]Node {node_id}:[/yellow] {errors}") - continue - - # Save output if requested - saved_path: Path | None = None - if output and result.images: - img_path = result.images[0] - img_data = get_image(str(img_path), url=url) - if img_data: - save_path = output if count == 1 else output.parent / f"{output.stem}_{i + 1:03d}{output.suffix}" - save_path.write_bytes(img_data) - saved_path = save_path - all_saved.append(save_path) - if not json_output: - console.print(f"[green]Saved:[/green] {save_path}") - elif not json_output: - console.print(f"[yellow]Could not download image: {img_path}[/yellow]") - - all_results.append( - { - "success": True, - "index": i, - "prompt_id": result.prompt_id, - "images": [str(img) for img in result.images], - "saved": str(saved_path) if saved_path else None, - } - ) + all_results.append( + { + "success": True, + "index": i, + "prompt_id": result.prompt_id, + "image": str(img_path), + "saved": str(saved_path) if saved_path else None, + } + ) if json_output: console.print_json( diff --git a/tensors/comfyui.py b/tensors/comfyui.py index c51357b..9db3f32 100644 --- a/tensors/comfyui.py +++ b/tensors/comfyui.py @@ -596,6 +596,7 @@ def _build_workflow( scheduler: str = "normal", lora_name: str | None = None, lora_strength: float = 1.0, + batch_size: int = 1, ) -> dict[str, Any]: """Build a text-to-image workflow from parameters. @@ -612,6 +613,7 @@ def _build_workflow( scheduler: Scheduler name lora_name: LoRA model filename (optional) lora_strength: LoRA strength (default 1.0) + batch_size: Number of images to generate in one workflow (default 1) Returns: ComfyUI workflow dict @@ -632,9 +634,10 @@ def _build_workflow( if model: workflow["4"]["inputs"]["ckpt_name"] = model - # Set dimensions + # Set dimensions and batch size workflow["5"]["inputs"]["width"] = width workflow["5"]["inputs"]["height"] = height + workflow["5"]["inputs"]["batch_size"] = batch_size # Set prompts workflow["6"]["inputs"]["text"] = prompt @@ -679,6 +682,7 @@ def generate_image( timeout: float = 600.0, lora_name: str | None = None, lora_strength: float = 1.0, + batch_size: int = 1, ) -> GenerationResult | None: """Generate an image using a simple text-to-image workflow. @@ -699,6 +703,7 @@ def generate_image( timeout: Maximum wait time in seconds lora_name: LoRA model filename (optional) lora_strength: LoRA strength (default 1.0) + batch_size: Number of images to generate in one workflow (default 1) Returns: GenerationResult with image paths, or None if generation failed @@ -731,6 +736,7 @@ def generate_image( scheduler=scheduler, lora_name=lora_name, lora_strength=lora_strength, + batch_size=batch_size, ) # Run workflow