Merge pull request #6 from NexaAI/xingyu/ci

[wip] add ci

.github/workflows/ci.yaml

@@ -0,0 +1,37 @@
+name: Python CI
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v3
+      with:
+        submodules: recursive  # This will clone the repository with all its submodules
+        fetch-depth: 0    # This fetches all history so you can access any version of the submodules
+
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'  # Specify the Python version you want
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install build pytest
+    - name: Build DLL
+      run: |
+        python -m pip install -e .
+    - name: Run tests
+      run: |
+        python -m pytest tests

nexa/gguf/nexa_inference_image.py

@@ -6,7 +6,6 @@
 import time
 from pathlib import Path
 
-from nexa.gguf.sd.stable_diffusion import StableDiffusion
 from nexa.general import pull_model
 from nexa.constants import (
     DEFAULT_IMG_GEN_PARAMS,
@@ -29,21 +28,22 @@ class NexaImageInference:
     A class used for loading image models and running image generation.
 
     Methods:
-    run_txt2img: Run the text-to-image generation loop.
-    run_img2img: Run the image-to-image generation loop.
-    run_streamlit: Run the Streamlit UI.
+        txt2img: (Used for SDK) Run the text-to-image generation loop.
+        img2img: (Used for SDK) Run the image-to-image generation loop.
+        run_streamlit: Run the Streamlit UI.
 
     Args:
-    model_path (str): Path or identifier for the model in Nexa Model Hub.
-    num_inference_steps (int): Number of inference steps.
-    width (int): Width of the output image.
-    height (int): Height of the output image.
-    guidance_scale (float): Guidance scale for diffusion.
-    output_path (str): Output path for the generated image.
-    random_seed (int): Random seed for image generation.
-    streamlit (bool): Run the inference in Streamlit UI.
+        model_path (str): Path or identifier for the model in Nexa Model Hub.
+        num_inference_steps (int): Number of inference steps.
+        width (int): Width of the output image.
+        height (int): Height of the output image.
+        guidance_scale (float): Guidance scale for diffusion.
+        output_path (str): Output path for the generated image.
+        random_seed (int): Random seed for image generation.
+        streamlit (bool): Run the inference in Streamlit UI.
 
     """
+
 
     def __init__(self, model_path, **kwargs):
         self.model_path = None
@@ -81,9 +81,10 @@ def __init__(self, model_path, **kwargs):
                 logging.error("Failed to load the model or pipeline.")
                 exit(1)
 
-    @SpinningCursorAnimation()
+    # @SpinningCursorAnimation()
     def _load_model(self, model_path: str):
         with suppress_stdout_stderr():
+            from nexa.gguf.sd.stable_diffusion import StableDiffusion
             self.model = StableDiffusion(
                 model_path=self.downloaded_path,
                 lora_model_dir=self.params.get("lora_dir", ""),
@@ -107,71 +108,124 @@ def _save_images(self, images):
             file_path = os.path.join(output_dir, file_name)
             image.save(file_path)
             logging.info(f"\nImage {i+1} saved to: {file_path}")
+
+    def txt2img(self, 
+                prompt, 
+                negative_prompt="",
+                cfg_scale=7.5,
+                width=512,
+                height=512,
+                sample_steps=20,
+                seed=0,
+                control_cond="",
+                control_strength=0.9):
+        """
+        Used for SDK. Generate images from text.
+
+        Args:
+            prompt (str): Prompt for the image generation.
+            negative_prompt (str): Negative prompt for the image generation.
 
-    def loop_txt2img(self):
+        Returns:
+            list: List of generated images.
+        """
+        images = self.model.txt_to_img(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            cfg_scale=cfg_scale,
+            width=width,
+            height=height,
+            sample_steps=sample_steps,
+            seed=seed,
+            control_cond=control_cond,
+            control_strength=control_strength,
+        )
+        return images
 
+    def run_txt2img(self):
         while True:
             try:
                 prompt = nexa_prompt("Enter your prompt: ")
                 negative_prompt = nexa_prompt(
                     "Enter your negative prompt (press Enter to skip): "
                 )
-                self._txt2img(prompt, negative_prompt)
+                try:
+                    images = self.txt2img(
+                        prompt, 
+                        negative_prompt,
+                        cfg_scale=self.params["guidance_scale"],
+                        width=self.params["width"],
+                        height=self.params["height"],
+                        sample_steps=self.params["num_inference_steps"],
+                        seed=self.params["random_seed"],
+                        control_cond=self.params.get("control_image_path", ""),
+                        control_strength=self.params.get("control_strength", 0.9),
+                    )
+                    self._save_images(images)
+                except Exception as e:
+                    logging.error(f"Error during text to image generation: {e}")
             except KeyboardInterrupt:
                 print(EXIT_REMINDER)
             except Exception as e:
                 logging.error(f"Error during generation: {e}", exc_info=True)
 
-    def _txt2img(self, prompt: str, negative_prompt: str):
+    def img2img(self, 
+                image_path, 
+                prompt, 
+                negative_prompt="",
+                cfg_scale=7.5,
+                width=512,
+                height=512,
+                sample_steps=20,
+                seed=0,
+                control_cond="",
+                control_strength=0.9):
         """
-        Generate images based on the given prompt, negative prompt, and parameters.
-        """
-        try:
-            images = self.model.txt_to_img(
-                prompt=prompt,
-                negative_prompt=negative_prompt if negative_prompt else "",
-                cfg_scale=self.params["guidance_scale"],
-                width=self.params["width"],
-                height=self.params["height"],
-                sample_steps=self.params["num_inference_steps"],
-                seed=self.params["random_seed"],
-                control_cond=self.params.get("control_image_path", ""),
-                control_strength=self.params.get("control_strength", 0.9),
-            )
-            self._save_images(images)
-        except Exception as e:
-            logging.error(f"Error during image generation: {e}")
+        Used for SDK. Generate images from an image.
 
-    def loop_img2img(self):
-        def _generate_images(image_path, prompt, negative_prompt):
-            """
-            Generate images based on the given prompt, negative prompt, and parameters.
-            """
-            try:
-                images = self.model.img_to_img(
-                    image=image_path,
-                    prompt=prompt,
-                    negative_prompt=negative_prompt if negative_prompt else "",
-                    cfg_scale=self.params["guidance_scale"],
-                    width=self.params["width"],
-                    height=self.params["height"],
-                    sample_steps=self.params["num_inference_steps"],
-                    seed=self.params["random_seed"],
-                    control_cond=self.params.get("control_image_path", ""),
-                    control_strength=self.params.get("control_strength", 0.9),
-                )
-                self._save_images(images)
-            except Exception as e:
-                logging.error(f"Error during image generation: {e}")
+        Args:
+            image_path (str): Path to the input image.
+            prompt (str): Prompt for the image generation.
+            negative_prompt (str): Negative prompt for the image generation.
 
+        Returns:
+            list: List of generated images.
+        """
+        images = self.model.img_to_img(
+            image=image_path,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            cfg_scale=cfg_scale,
+            width=width,
+            height=height,
+            sample_steps=sample_steps,
+            seed=seed,
+            control_cond=control_cond,
+            control_strength=control_strength,
+        )
+        return images
+
+    def run_img2img(self):
         while True:
             try:
                 image_path = nexa_prompt("Enter the path to your image: ")
                 prompt = nexa_prompt("Enter your prompt: ")
                 negative_prompt = nexa_prompt(
                     "Enter your negative prompt (press Enter to skip): "
                 )
-                _generate_images(image_path, prompt, negative_prompt)
+                images = self.img2img(image_path, 
+                                      prompt, 
+                                      negative_prompt,
+                                      cfg_scale=self.params["guidance_scale"],
+                                      width=self.params["width"],
+                                      height=self.params["height"],
+                                      sample_steps=self.params["num_inference_steps"],
+                                      seed=self.params["random_seed"],
+                                      control_cond=self.params.get("control_image_path", ""),
+                                        control_strength=self.params.get("control_strength", 0.9),
+                                    )
+
+                self._save_images(images)
             except KeyboardInterrupt:
                 print(EXIT_REMINDER)
             except Exception as e:
@@ -257,6 +311,6 @@ def run_streamlit(self, model_path: str):
         inference.run_streamlit(model_path)
     else:
         if args.img2img:
-            inference.loop_img2img()
+            inference.run_img2img()
         else:
-            inference.loop_txt2img()
+            inference.run_txt2img()