Skip to content

Commit

Permalink
Merge pull request #7 from rtuszik/dev
Browse files Browse the repository at this point in the history
Various UI Changes, Docker Fixes, Build Process Changes
  • Loading branch information
rtuszik authored Aug 26, 2024
2 parents 8d38712 + d77aab6 commit 7498c3d
Show file tree
Hide file tree
Showing 5 changed files with 97 additions and 43 deletions.
7 changes: 7 additions & 0 deletions .github/workflows/docker-build-push.yml
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@ on:
push:
branches:
- main
paths:
- src/**
- Dockerfile

env:
REGISTRY: ghcr.io
Expand Down Expand Up @@ -35,6 +38,8 @@ jobs:
tags: |
type=raw,value=latest
type=sha,prefix={{branch}}-
type=ref,event=branch
- name: Build and push Docker image
uses: docker/build-push-action@v4
Expand All @@ -43,3 +48,5 @@ jobs:
push: true
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha
cache-to: type=gha,mode=max
21 changes: 10 additions & 11 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,17 +22,16 @@ Docker is the recommended way to run this application. It ensures consistent env
### Docker Compose Setup
1. Create a `docker-compose.yml` file with the following content:
```yaml
services:
replicate-flux-lora:
image: ghcr.io/rtuszik/replicate-flux-lora:latest
container_name: replicate-flux-lora
environment:
- REPLICATE_API_TOKEN=${REPLICATE_API_TOKEN}
ports:
- "8080:8080"
volumes:
- ${HOST_OUTPUT_DIR}:/app/output
restart: unless-stopped
services:
replicate-flux-lora:
image: ghcr.io/rtuszik/replicate-flux-lora:latest
container_name: replicate-flux-lora
env_file: .env
ports:
- "8080:8080"
volumes:
- ${HOST_OUTPUT_DIR}:/app/output
restart: unless-stopped
```
2. Create a `.env` file in the same directory with the following content:
Expand Down
3 changes: 1 addition & 2 deletions docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@ services:
replicate-flux-lora:
image: ghcr.io/rtuszik/replicate-flux-lora:latest
container_name: replicate-flux-lora
environment:
- REPLICATE_API_TOKEN=${REPLICATE_API_TOKEN}
env_file: .env
ports:
- "8080:8080"
volumes:
Expand Down
96 changes: 73 additions & 23 deletions src/gui.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,7 @@
from loguru import logger
from nicegui import events, ui

# Configure Loguru
logger.remove() # Remove the default handler
logger.remove()
logger.add(
sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO"
)
Expand Down Expand Up @@ -57,36 +56,67 @@ def __init__(self, image_generator):
self.load_settings()
self.flux_fine_tune_models = self.image_generator.get_flux_fine_tune_models()
self.user_added_models = self.settings.get("user_added_models", [])

self._attributes = [
"prompt",
"flux_model",
"aspect_ratio",
"num_outputs",
"lora_scale",
"num_inference_steps",
"guidance_scale",
"output_format",
"output_quality",
"disable_safety_checker",
"width",
"height",
"seed",
]

for attr in self._attributes:
setattr(self, attr, self.settings.get(attr, None))

self.setup_ui()
logger.info("ImageGeneratorGUI initialized")

def __getattr__(self, name):
if name in self._attributes:
return self.__dict__.get(name, None)
return super().__getattribute__(name)

def setup_ui(self):
ui.dark_mode().enable()

with ui.column().classes("w-full max-w-full mx-auto p-4 space-y-4"):
with ui.column().classes("w-full max-w-full mx-auto space-y-8"):
with ui.card().classes("w-full"):
ui.label("Image Generator").classes("text-2xl font-bold mb-4")
with ui.row().classes("w-full"):
with ui.column().classes("w-1/2 pr-2"):
ui.label("Flux LoRA API").classes("text-2xl font-bold mb-4")
with ui.row():
with ui.column(wrap=False):
self.setup_left_panel()
with ui.column().classes("w-1/2 pl-2"):
with ui.column(wrap=False):
self.setup_right_panel()
ui.separator()
self.setup_bottom_panel()
logger.info("UI setup completed")

def setup_left_panel(self):
self.replicate_model_input = ui.input(
"Replicate Model", value=self.settings.get("replicate_model", "")
).classes("w-full")
self.replicate_model_input = (
ui.input("Replicate Model", value=self.settings.get("replicate_model", ""))
.classes("w-full")
.tooltip("Enter the Replicate model URL or identifier")
)
self.replicate_model_input.on("change", self.update_replicate_model)

self.flux_models_select = ui.select(
options=self.flux_fine_tune_models,
label="Flux Fine-Tune Models",
value=None,
on_change=self.select_flux_model,
).classes("w-full")

self.flux_models_select = (
ui.select(
options=self.flux_fine_tune_models,
label="Flux Fine-Tune Models",
value=None,
on_change=self.select_flux_model,
)
.classes("w-full")
.tooltip("Select Model")
)
with ui.row().classes("w-full"):
self.new_model_input = ui.input(label="New Model").classes("w-3/4")
ui.button("Add Model", on_click=self.add_user_model).classes("w-1/4")
Expand Down Expand Up @@ -121,6 +151,9 @@ def setup_left_panel(self):
value=self.settings.get("flux_model", "dev"),
)
.classes("w-full")
.tooltip(
"Which model to run inferences with. the dev model needs around 28 steps but the schnell model only needs around 4 steps."
)
.bind_value(self, "flux_model")
)

Expand All @@ -145,6 +178,9 @@ def setup_left_panel(self):
)
.classes("w-full")
.bind_value(self, "aspect_ratio")
.tooltip(
"Width of the generated image. Optional, only used when aspect_ratio=custom. Must be a multiple of 16 (if it's not, it will be rounded to nearest multiple of 16)"
)
)
self.aspect_ratio_select.on("change", self.toggle_custom_dimensions)

Expand All @@ -157,13 +193,19 @@ def setup_left_panel(self):
)
.classes("w-full")
.bind_value(self, "width")
.tooltip(
"Width of the generated image. Optional, only used when aspect_ratio=custom. Must be a multiple of 16 (if it's not, it will be rounded to nearest multiple of 16)"
)
)
self.height_input = (
ui.number(
"Height", value=self.settings.get("height", 1024), min=256, max=1440
)
.classes("w-full")
.bind_value(self, "height")
.tooltip(
"Height of the generated image. Optional, only used when aspect_ratio=custom. Must be a multiple of 16 (if it's not, it will be rounded to nearest multiple of 16)"
)
)

self.num_outputs_input = (
Expand All @@ -172,6 +214,7 @@ def setup_left_panel(self):
)
.classes("w-full")
.bind_value(self, "num_outputs")
.tooltip("Number of images to output.")
)
self.lora_scale_input = (
ui.number(
Expand All @@ -182,6 +225,9 @@ def setup_left_panel(self):
step=0.1,
)
.classes("w-full")
.tooltip(
"Determines how strongly the LoRA should be applied. Sane results between 0 and 1."
)
.bind_value(self, "lora_scale")
)
self.num_inference_steps_input = (
Expand All @@ -192,6 +238,7 @@ def setup_left_panel(self):
max=50,
)
.classes("w-full")
.tooltip("Number of Inference Steps")
.bind_value(self, "num_inference_steps")
)
self.guidance_scale_input = (
Expand All @@ -203,6 +250,7 @@ def setup_left_panel(self):
step=0.1,
)
.classes("w-full")
.tooltip("Guidance Scale for the diffusion process")
.bind_value(self, "guidance_scale")
)
self.seed_input = (
Expand All @@ -222,6 +270,7 @@ def setup_left_panel(self):
value=self.settings.get("output_format", "webp"),
)
.classes("w-full")
.tooltip("Format of the output images")
.bind_value(self, "output_format")
)
self.output_quality_input = (
Expand All @@ -232,14 +281,18 @@ def setup_left_panel(self):
max=100,
)
.classes("w-full")
.tooltip(
"Quality when saving the output images, from 0 to 100. 100 is best quality, 0 is lowest quality. Not relevant for .png outputs"
)
.bind_value(self, "output_quality")
)
self.disable_safety_checker_switch = (
ui.switch(
"Disable Safety Checker",
value=self.settings.get("disable_safety_checker", False),
value=self.settings.get("disable_safety_checker", True),
)
.classes("w-full")
.tooltip("Disable safety checker for generated images.")
.bind_value(self, "disable_safety_checker")
)

Expand Down Expand Up @@ -288,7 +341,6 @@ def setup_right_panel(self):
self.spinner = ui.spinner(type="infinity", size="xl")
self.spinner.visible = False

# Add gallery view
self.gallery_container = ui.column().classes("w-full mt-4")
self.lightbox = Lightbox()

Expand All @@ -297,6 +349,7 @@ def setup_bottom_panel(self):
ui.textarea("Prompt", value=self.settings.get("prompt", ""))
.classes("w-full")
.bind_value(self, "prompt")
.props("clearable")
)
self.generate_button = ui.button(
"Generate Images", on_click=self.start_generation
Expand Down Expand Up @@ -342,7 +395,6 @@ async def start_generation(self):
)
return

# Ensure the model is set in the ImageGenerator
self.image_generator.set_model(self.replicate_model_input.value)

self.save_settings()
Expand Down Expand Up @@ -392,9 +444,7 @@ async def download_and_display_images(self, image_urls):
response = await client.get(url)
if response.status_code == 200:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
url_part = urllib.parse.urlparse(url).path.split("/")[-2][
:8
] # Get first 8 chars of the unique part
url_part = urllib.parse.urlparse(url).path.split("/")[-2][:8]
file_name = f"generated_image_{timestamp}_{url_part}_{i+1}.png"
file_path = Path(self.folder_path) / file_name
with open(file_path, "wb") as f:
Expand Down
13 changes: 6 additions & 7 deletions src/main.py
Original file line number Diff line number Diff line change
@@ -1,24 +1,24 @@
import sys

from gui import create_gui
from loguru import logger
from nicegui import ui

from gui import create_gui
from replicate_api import ImageGenerator

# Configure Loguru
logger.remove() # Remove the default handler
logger.remove()
logger.add(
sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO"
)
logger.add(
"main.log", rotation="10 MB", format="{time} {level} {message}", level="INFO"
)

# Create the ImageGenerator instance

logger.info("Initializing ImageGenerator")
generator = ImageGenerator()

# Create and setup the GUI

logger.info("Creating and setting up GUI")


Expand All @@ -28,7 +28,6 @@ async def main_page():
logger.info("NiceGUI server is running")


# Run the NiceGUI server
logger.info("Starting NiceGUI server")
# ui.run(port=8080)

ui.run(title="Replicate Flux LoRA", port=8080)

0 comments on commit 7498c3d

Please sign in to comment.