feat: create separate image for dev and prod

Dockerfile

@@ -1,10 +1,14 @@
-# Use RunPod's base image
 FROM runpod/base:0.6.1-cuda12.2.0
 
-# Only install absolutely necessary packages for your project
-RUN apt-get update && apt-get install -y ffmpeg aria2 git unzip
+# Environment variable to differentiate between development and production
+ARG ENVIRONMENT=development
+ENV ENVIRONMENT=${ENVIRONMENT}
 
-# Clone the ComfyUI repository
+# Install necessary system packages
+RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg aria2 git unzip && \
+    rm -rf /var/lib/apt/lists/*
+
+# Clone the ComfyUI repository 
 RUN git clone https://github.com/comfyanonymous/ComfyUI /content/ComfyUI
 
 # Download required model files
@@ -18,8 +22,21 @@ RUN mkdir -p /content/ComfyUI/models/unet && \
     mkdir -p /content/ComfyUI/models/loras && \
     aria2c --console-log-level=error -c -x 16 -s 16 -k 1M "https://civitai.com/api/download/models/896422?type=Model&format=SafeTensor" -d /content/ComfyUI/models/loras -o zanshou-kin-flux-ueno-manga-style.safetensors
 
-# Reset the working directory to the base image's root
-WORKDIR /
+# Copy requirements.txt earlier so it can be used for installing dependencies
+COPY builder/requirements.txt /requirements.txt
+
+# Install Python dependencies if building for production
+RUN if [ "$ENVIRONMENT" = "production" ]; then \
+    python3.10 -m pip install --upgrade pip && \
+    python3.10 -m pip install --no-cache-dir -r /requirements.txt; \
+    fi
+
+# Copy all files from src to the root directory
+COPY ./src/ /
 
-# Keep the base image's entrypoint
-CMD ["/start.sh"]
+# Final CMD based on environment
+CMD if [ "$ENVIRONMENT" = "development" ]; then \
+    /start.sh; \
+    else \
+    python3.10 -u /handler.py; \
+    fi

builder/requirements.txt

@@ -0,0 +1,24 @@
+# Required Python packages get listed here, one per line.
+# Lock the version number to avoid unexpected changes.
+
+# You can also install packages from a git repository, e.g.:
+# git+https://github.com/runpod/runpod-python.git
+# To learn more, see https://pip.pypa.io/en/stable/reference/requirements-file-format/
+
+torch
+torchvision
+torchaudio
+opencv-python
+imageio
+imageio-ffmpeg
+ffmpeg-python
+av
+runpod
+xformers==0.0.25
+torchsde==0.2.6
+einops==0.8.0
+diffusers==0.28.0
+transformers==4.41.2
+accelerate==0.30.1
+numpy<2
+

docker-bake.hcl

@@ -0,0 +1,33 @@
+variable "DOCKERHUB_REPO" {
+  default = "timpietruskyblibla"
+}
+
+variable "DOCKERHUB_IMG" {
+  default = "flux-1-dev-lora-worker"
+}
+
+variable "RELEASE_VERSION" {
+  default = "latest"
+}
+
+group "default" {
+  targets = ["development", "production"]
+}
+
+target "development" {
+  context = "."
+  dockerfile = "Dockerfile"
+  args = {
+    ENVIRONMENT = "development"
+  }
+  tags = ["${DOCKERHUB_REPO}/${DOCKERHUB_IMG}:dev"]
+}
+
+target "production" {
+  context = "."
+  dockerfile = "Dockerfile"
+  args = {
+    ENVIRONMENT = "production"
+  }
+  tags = ["${DOCKERHUB_REPO}/${DOCKERHUB_IMG}:latest"]
+}

model.json

@@ -1,13 +1,16 @@
 {
-    "positive_prompt": "A cat holding a sign that says 'I LOVE TOST!'",
-    "seed": 0,
-    "steps": 20,
+  "input": {
+    "positive_prompt": "a black and white drawing of a city",
+    "width": 512,
+    "height": 512,
+    "seed": 13379,
+    "steps": 24,
     "guidance": 3.5,
-    "lora_url": "https://huggingface.co/TostAI/flux-1-dev-lora/resolve/main/8iiu0y.safetensors",
-    "lora_strength_model": 1,
-    "lora_strength_clip": 1,
+    "lora_strength_model": 0.8,
+    "lora_strength_clip": 0.8,
     "sampler_name": "euler",
     "scheduler": "simple",
-    "width": 1024,
-    "height": 1024
-}
+    "job_id": "test-job-123",
+    "lora_name": "zanshou-kin-flux-ueno-manga-style.safetensors"
+  }
+}

flux1-dev-lora-worker-project/src/handler.py → src/handler.py

@@ -1,6 +1,4 @@
 import os
-import json
-import requests
 import runpod
 import random
 import time
@@ -13,11 +11,9 @@
 
 sys.path.append("/content/ComfyUI")
 
-import nodes
 from nodes import NODE_CLASS_MAPPINGS
 from comfy_extras import nodes_custom_sampler
 from comfy_extras import nodes_flux
-from comfy import model_management
 
 # Initialize Model Loaders
 DualCLIPLoader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
@@ -64,17 +60,18 @@ def generate(input):
     width = values.get("width", 512)
     height = values.get("height", 512)
     seed = values.get("seed", 0)
-    steps = values.get("steps", 50)
-    guidance = values.get("guidance", 7.5)
+    steps = values.get("steps", 25)
+    guidance = values.get("guidance", 3.5)
     lora_strength_model = values.get("lora_strength_model", 0.8)
     lora_strength_clip = values.get("lora_strength_clip", 0.8)
-    sampler_name = values.get("sampler_name", "Euler")
-    scheduler = values.get("scheduler", "default")
+    sampler_name = values.get("sampler_name", "euler")
+    scheduler = values.get("scheduler", "simple")
     job_id = values.get("job_id", "test-job-123")
     lora_name = values.get("lora_name", "zanshou-kin-flux-ueno-manga-style.safetensors")
+    image_path = None
 
     # Path to the LoRA model based on lora_name
-    lora_file_path = f"models/loras/{lora_name}"
+    lora_file_path = f"/content/ComfyUI/models/loras/{lora_name}"
 
     # Validate if the specified LoRA model exists
     if not os.path.exists(lora_file_path):
@@ -95,9 +92,10 @@ def generate(input):
     print(f"Using seed: {seed}")
 
     try:
+
         # Load LoRA models from the specified file
         unet_lora, clip_lora = LoraLoader.load_lora(
-            unet, clip, lora_file_path, lora_strength_model, lora_strength_clip
+            unet, clip, lora_name, lora_strength_model, lora_strength_clip
         )
 
         # Encode the positive prompt
@@ -151,7 +149,7 @@ def generate(input):
 
     finally:
         # Clean up the generated image file
-        if os.path.exists(image_path):
+        if image_path is not None and os.path.exists(image_path):
             os.remove(image_path)