Rework - Adding parallel exec option, fixing unit tests, improving co…

…de, updating readme

README.md

@@ -4,10 +4,11 @@
 [![codecov](https://codecov.io/gh/leopedroso45/Stable-Diffusion-ImageGen/branch/main/graph/badge.svg?token=YOUR_TOKEN)](https://codecov.io/gh/leopedroso45/Stable-Diffusion-ImageGen)
 ![License](https://img.shields.io/github/license/leopedroso45/Stable-Diffusion-ImageGen)
 
-`sevsd` is a Python package designed to simplify the integration of Stable Diffusion image generation into various applications. Utilizing Hugging Face's `diffusers` library, `sevsd` provides a straightforward and flexible interface for generating images based on textual prompts. Whether you're building HTTP APIs, high-level services, or other applications, `sevsd` streamlines the process of incorporating AI-driven image generation.
+`sevsd` is a Python package specifically designed to make the process of generating images using Stable Diffusion models as simple as possible. The package enables image generation with just a single function call, greatly simplifying the integration of Stable Diffusion into various applications. Utilizing Hugging Face's `diffusers` library, `sevsd` provides an intuitive and flexible interface for generating images based on textual prompts. This makes it an ideal choice for building HTTP APIs, high-level services, or any application requiring AI-driven image generation.
 
 ## Features
 
+- Simplified interface for Stable Diffusion image generation, enabling the creation of images with just a single function call.
 - Easy integration of Stable Diffusion model into Python applications.
 - Customizable image generation based on user-defined tasks and configurations.
 - Batch processing capabilities for handling multiple tasks efficiently.
@@ -37,31 +38,64 @@ Import and use `sevsd` in your Python project:
 ```python
 from sevsd import do_work
 
-# Define your configuration and tasks
-configs = [("CompVis/stable-diffusion-v1-4", "./model_cache")]
-tasks = [("A scenic landscape", None, 50, 1, 7.5)]
-
-# Process tasks
-do_work(configs, tasks, "./generated-images")
+# Define your models and jobs
+models = [
+    {
+        "name": './model_cache/model1.safetensors',
+        "executor": {
+            "labels": [1],
+            "num_of_exec": 1,
+            "cfg_scale": 7,
+            "inference_steps": 100,
+        }
+    },
+    {
+        "name": './model_cache/model2.safetensors',
+        "executor": {
+            "labels": [2],
+            "num_of_exec": 2,
+            "cfg_scale": 6,
+            "inference_steps": 50,
+        }
+    },
+]
+
+jobs = [
+    {
+        "label": 1,
+        "prompt": 'A scenic landscape',
+        "negative_prompt": "blurred image, black and white, watermarked image",
+    },
+    {
+        "label": 2,
+        "prompt": 'A person wearing a mask',
+        "negative_prompt": 'deformed anatomy, hand-drawn image, blurred image',
+    },
+]
+
+do_work(models, jobs, './generated-images')
 ```
 
-This example demonstrates a basic usage scenario. Customize the `configs` and `tasks` as needed for your application.
+This example demonstrates a basic usage scenario. Customize the `models` and `jobs` as needed for your application.
 
 ## Components
 
-- `setup_pipeline`: Initializes the Stable Diffusion pipeline with given configurations.
-- `process_task`: Processes the list of tasks, generating and saving images.
-- `generate_image`: Handles the image generation process for each task.
+- `setup_pipeline`: Prepares the Stable Diffusion pipeline with the specified model configuration.
+- `process_task`: Processes individual tasks, generating and saving images based on job specifications.
+- `generate_image`: Handles the image generation process for each job.
 - `setup_device`: Sets up the computation device (GPU or CPU) for image generation.
+- `check_os_path`: Ensures the output path exists or creates it.
+- `check_cuda_and_clear_cache`: Manages GPU memory and cache for efficient processing.
+- `do_work`: Central function to orchestrate the processing of jobs with corresponding models.
 
 ## Customization
 
-You can customize the image generation process by modifying the `tasks` list with different prompts, inference steps, and other parameters. The `configs` list allows for different model configurations, enabling diverse image styles.
+You can customize the image generation process by adjusting the `models` and `jobs` lists. Define different prompts, model paths, execution parameters, and more to cater to diverse image styles and requirements.
 
 ## Note
 
 - Ensure sufficient GPU memory if using CUDA.
-- The package is optimized for batch processing. Modify `tasks` and `configs` to fit your requirements.
+- The package is optimized for flexible handling of various job and model configurations.
 - For detailed examples and advanced usage, refer to the source code documentation.
 
 ## Contributing

sevsd/do_work.py

@@ -1,15 +1,50 @@
 from sevsd.setup_pipeline import setup_pipeline
 from sevsd.process_task import process_task
 
-def do_work(configs, tasks, path, parallel_exec=True, **kwargs):
-    task_dict = {task["task_id"]: [] for task in tasks}
-    for task in tasks:
-        task_dict[task["task_id"]].append(task)
-
-    for config in configs:
-        pipeline = setup_pipeline(config["model_info"], **kwargs)
-        task_ids = config.get("task_ids", [])
-        for task_id in task_ids:
-            if task_id in task_dict:
-                for task in task_dict[task_id]:
-                    process_task(task["details"], pipeline, path, parallel_exec)
+models = [
+    {
+        "name": "./model_cache/model1.safetensors",
+        "executor": {
+            "labels": [1],
+            "num_of_exec": 1,
+            "cfg_scale": 7,
+            "inference_steps": 100,
+        }
+    },
+    {
+        "name": "./model_cache/model2.safetensors",
+        "executor": {
+            "labels": [2],
+            "num_of_exec": 2,
+            "cfg_scale": 6,
+            "inference_steps": 50,
+        }
+    },
+]
+
+jobs = [
+    {
+        "label": 1,
+        "prompt": "A scenic landscape",
+        "negative_prompt": "blurred image, black and white, watermarked image",
+    },
+    {
+        "label": 2,
+        "prompt": "A person wearing a mask",
+        "negative_prompt": "deformed anatomy, hand-drawn image, blurred image",
+    },
+]
+
+def do_work(models, jobs, image_path, parallel_exec=True, **kwargs):
+    job_dict = {job['label']: [] for job in jobs}
+    for job in jobs:
+        job_dict[job['label']].append(job)
+
+    for model in models:
+        pipeline = setup_pipeline(model["name"], **kwargs)
+        labels = model.get("executor", {}).get("labels", [])
+        for label in labels:
+            if label in job_dict:
+                for job in job_dict[label]:
+                    executor = model.get("executor", {})
+                    process_task(job, pipeline, executor, image_path, parallel_exec)

sevsd/generate_image.py

@@ -1,14 +1,14 @@
 import torch
-
-def check_and_clear_cache():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
 
-def generate_image(task_details, pipeline, parallel_exec=True, **kwargs):
-    prompt, negative_prompt, num_inference_steps, num_images, cfg = task_details
+def generate_image(job, pipeline, executor, parallel_exec=True, **kwargs):
+
+    prompt = job.get("prompt")
+    negative_prompt = job.get("negative_prompt")
+    num_inference_steps = executor.get("inference_steps")
+    num_images = executor.get("num_of_exec")
+    cfg = executor.get("cfg_scale")
 
     def execute_pipeline(num_images):
-        check_and_clear_cache()
         return pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -17,17 +17,14 @@ def execute_pipeline(num_images):
             guidance_scale=cfg,
             **kwargs
         )
-
     try:
         with torch.no_grad():
             if parallel_exec:
                 output = execute_pipeline(num_images)
                 return output["images"]
             else:
-                output_list = [execute_pipeline(1)["images"][0] for _ in range(num_images)]
-                return output_list
+                output = execute_pipeline(1)
+                return output["images"]
     except RuntimeError as e:
         print(f"Runtime error: {e}")
-        return None
-    finally:
-        check_and_clear_cache()
+        return None

sevsd/process_task.py

@@ -6,23 +6,31 @@
 from sevsd.generate_image import generate_image
 import torch
 
-def process_task(tasks, pipeline, path, parallel_exec=True):
+def process_task(job, pipeline, executor, path, parallel_exec=True):
+
+    def call_generate_image():
+        images = generate_image(job, pipeline, executor, parallel_exec)
+        if images is not None:
+            for image in images:
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S%f")
+                image_path = f"{path}/generated_image_{timestamp}.png"
+                image.save(image_path)
+                print(f"[sevsd] - image saved at {image_path}")
+        else:
+            print("[sevsd] - image generation failed due to memory constraints.")
+        check_cuda_and_clear_cache()
+
     try:
         path = check_os_path(path)
-        if tasks is not None:
-            for task in tasks:
-                images = generate_image(task, pipeline, parallel_exec)
-                if images is not None:
-                    for image in images:
-                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S%f")
-                        image_path = f"{path}/generated_image_{timestamp}.png"
-                        image.save(image_path)
-                        print(f"Image saved at {image_path}")
-                else:
-                    print("Image generation failed due to memory constraints.")
-                check_cuda_and_clear_cache()
+        if job is not None:
+            if parallel_exec is not True:
+                num_images = executor.get("num_of_exec", 1)
+                for _ in range(num_images):
+                    call_generate_image()
+            else:
+                call_generate_image()
     except Exception as e:
-        print(f"Exception: {e}")
+        print(f"[sevsd] - exception: {e}")
     finally:
         check_cuda_and_clear_cache()
 
@@ -35,5 +43,5 @@ def check_cuda_and_clear_cache():
 def check_os_path(path):
     if not os.path.exists(path):
         os.makedirs(path)
-        print(f"Created path: {path}")
+        print(f"[sevsd] - created path: {path}")
     return path

sevsd/setup_pipeline.py

@@ -1,26 +1,25 @@
 from sevsd.setup_device import setup_device
 from diffusers import StableDiffusionPipeline
 
-def setup_pipeline(config, **kwargs):
+def setup_pipeline(pretrained_model_link_or_path, **kwargs):
     device = setup_device()
-    pretrained_model_link_or_path, cache_dir = config
 
     default_kwargs = {
         "use_safetensors": False,
         "load_safety_checker": False,
         "requires_safety_checker": False,
-        "cache_dir": cache_dir,
     }
 
-    default_kwargs.update(kwargs)
-
     if pretrained_model_link_or_path.endswith(".safetensors"):
         default_kwargs["use_safetensors"] = True
+        default_kwargs.update(kwargs)
+
         pipeline = StableDiffusionPipeline.from_single_file(
             pretrained_model_link_or_path,
             **default_kwargs
         )
     else:
+        default_kwargs.update(kwargs)
         pipeline = StableDiffusionPipeline.from_pretrained(
             pretrained_model_link_or_path,
             **default_kwargs

tests/test_do_work.py

@@ -12,20 +12,26 @@ def test_do_work(self, mock_setup_pipeline, mock_process_task):
         mock_pipeline = MagicMock()
         mock_setup_pipeline.return_value = mock_pipeline
 
-        fake_configs = [
-            {"model_info": ("model_path", "cache_path"), "task_ids": [1, 2]}
+        fake_models = [
+            {
+                "name": "model_path",
+                "executor": {
+                    "labels": [1, 2],
+                    "num_of_exec": 2,
+                    "cfg_scale": 7
+                }
+            }
         ]
-        fake_tasks = [
-            {"task_id": 1, "details": ("prompt1", None, 50, 1, 7.5)},
-            {"task_id": 2, "details": ("prompt2", None, 30, 2, 8.0)}
+        fake_jobs = [
+            {"label": 1, "details": ("prompt1", None, 50, 1, 7.5)},
+            {"label": 2, "details": ("prompt2", None, 30, 2, 8.0)}
         ]
         fake_path = "test_path"
 
-        do_work(fake_configs, fake_tasks, fake_path)
+        do_work(fake_models, fake_jobs, fake_path)
 
-        mock_setup_pipeline.assert_called_once_with(('model_path', 'cache_path'))
-        mock_process_task.assert_any_call(("prompt1", None, 50, 1, 7.5), mock_pipeline, fake_path, True)
-        mock_process_task.assert_any_call(("prompt2", None, 30, 2, 8.0), mock_pipeline, fake_path, True)
+        mock_process_task.assert_any_call(fake_jobs[0], mock_pipeline, fake_models[0]['executor'], fake_path, True)
+        mock_process_task.assert_any_call(fake_jobs[1], mock_pipeline, fake_models[0]['executor'], fake_path, True)
 
 if __name__ == '__main__':
     unittest.main()

tests/test_generate_image.py

@@ -8,41 +8,44 @@ class TestGenerateImage(unittest.TestCase):
 
     @patch('sevsd.generate_image.torch')
     def test_generate_image(self, mock_torch):
-        fake_args = ("prompt", None, 50, 1, 7.5)
+        fake_job = {"prompt": "prompt", "negative_prompt": None}
+        fake_executor = {"inference_steps": 50, "num_of_exec": 1, "cfg_scale": 7.5}
         fake_pipeline = MagicMock()
         mock_output = {'images': [MagicMock()]}
         fake_pipeline.return_value = mock_output
 
-        images = generate_image(fake_args, fake_pipeline, parallel_exec=True)
+        images = generate_image(fake_job, fake_pipeline, fake_executor, parallel_exec=True)
 
         self.assertEqual(len(images), len(mock_output['images']))
-        mock_torch.cuda.empty_cache.assert_called()
-    
+        mock_torch.no_grad.assert_called()
+
     @patch('sevsd.generate_image.torch')
     def test_generate_image_sequential_execution(self, mock_torch):
-        fake_args = ("prompt", None, 50, 10, 7.5)  # num_images = 10 para testar a execução sequencial
+        fake_job = {"prompt": "prompt", "negative_prompt": None}
+        fake_executor = {"inference_steps": 50, "num_of_exec": 10, "cfg_scale": 7.5}  # num_images = 10 para testar a execução sequencial
         fake_pipeline = MagicMock()
         fake_image = MagicMock()
-        mock_output = {'images': [fake_image]}
+        mock_output = {'images': [fake_image] * 10}
         fake_pipeline.return_value = mock_output
 
-        images = generate_image(fake_args, fake_pipeline, parallel_exec=False)
+        images = generate_image(fake_job, fake_pipeline, fake_executor, parallel_exec=False)
 
         self.assertEqual(len(images), 10)
         self.assertEqual(images[0], fake_image)
-        mock_torch.cuda.empty_cache.assert_called()
+        mock_torch.no_grad.assert_called()
 
     @patch('sevsd.generate_image.torch')
     def test_generate_image_runtime_error(self, mock_torch):
-        fake_args = ("prompt", None, 50, 1, 7.5)
+        fake_job = {"prompt": "prompt", "negative_prompt": None}
+        fake_executor = {"inference_steps": 50, "num_of_exec": 1, "cfg_scale": 7.5}
         fake_pipeline = MagicMock()
 
         fake_pipeline.side_effect = RuntimeError("Test error")
 
-        images = generate_image(fake_args, fake_pipeline, parallel_exec=False)
+        images = generate_image(fake_job, fake_pipeline, fake_executor, parallel_exec=False)
 
         self.assertIsNone(images)
-        mock_torch.cuda.empty_cache.assert_called()
+        mock_torch.no_grad.assert_called()
 
 if __name__ == '__main__':
     unittest.main()