Compare commits

...

7 commits

Author SHA1 Message Date
dependabot[bot] 6ca8dc359f
Merge eb543cbe3a into 6291e72129 2025-12-02 17:57:35 +00:00
oobabooga 6291e72129 Remove quanto for now (requires messy compilation) 2025-12-02 09:57:18 -08:00
oobabooga b3666e140d
Add image generation support (#7328) 2025-12-02 14:55:38 -03:00
oobabooga a83821e941 Revert "UI: Optimize typing in all textareas"
This reverts commit e24ba92ef2.
2025-12-01 10:34:23 -08:00
oobabooga 24fd963c38 Merge remote-tracking branch 'refs/remotes/origin/dev' into dev 2025-12-01 08:06:08 -08:00
oobabooga e24ba92ef2 UI: Optimize typing in all textareas 2025-12-01 08:05:21 -08:00
dependabot[bot] eb543cbe3a
Update safetensors requirement in /requirements/full
Updates the requirements on [safetensors](https://github.com/huggingface/safetensors) to permit the latest version.
- [Release notes](https://github.com/huggingface/safetensors/releases)
- [Changelog](https://github.com/huggingface/safetensors/blob/main/RELEASE.md)
- [Commits](https://github.com/huggingface/safetensors/compare/v0.6.0-rc0...v0.7.0)

---
updated-dependencies:
- dependency-name: safetensors
  dependency-version: 0.7.0
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
2025-11-24 20:38:05 +00:00
18 changed files with 1290 additions and 23 deletions

View file

@ -28,6 +28,8 @@ A Gradio web UI for Large Language Models.
- 100% offline and private, with zero telemetry, external resources, or remote update requests.
- **File attachments**: Upload text files, PDF documents, and .docx documents to talk about their contents.
- **Vision (multimodal models)**: Attach images to messages for visual understanding ([tutorial](https://github.com/oobabooga/text-generation-webui/wiki/Multimodal-Tutorial)).
Image generation: A dedicated tab for diffusers models like Z-Image-Turbo and Qwen-Image. Features 4-bit/8-bit quantization and a persistent gallery with metadata (tutorial).
- **Image generation**: A dedicated tab for `diffusers` models like **Z-Image-Turbo** and **Qwen-Image**. Features 4-bit/8-bit quantization and a persistent gallery with metadata ([tutorial](https://github.com/oobabooga/text-generation-webui/wiki/Image-Generation-Tutorial)).
- **Web search**: Optionally search the internet with LLM-generated queries to add context to the conversation.
- Aesthetic UI with dark and light themes.
- Syntax highlighting for code blocks and LaTeX rendering for mathematical expressions.
@ -432,6 +434,7 @@ https://colab.research.google.com/github/oobabooga/text-generation-webui/blob/ma
https://www.reddit.com/r/Oobabooga/
## Acknowledgment
## Acknowledgments
In August 2023, [Andreessen Horowitz](https://a16z.com/) (a16z) provided a generous grant to encourage and support my independent work on this project. I am **extremely** grateful for their trust and recognition.
- In August 2023, [Andreessen Horowitz](https://a16z.com/) (a16z) provided a generous grant to encourage and support my independent work on this project. I am **extremely** grateful for their trust and recognition.
- This project was inspired by [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) and wouldn't exist without it.

View file

@ -93,11 +93,11 @@ ol li p, ul li p {
display: inline-block;
}
#notebook-parent-tab, #chat-tab, #parameters, #chat-settings, #lora, #training-tab, #model-tab, #session-tab, #character-tab {
#notebook-parent-tab, #chat-tab, #parameters, #chat-settings, #lora, #training-tab, #model-tab, #session-tab, #character-tab, #image-ai-tab {
border: 0;
}
#notebook-parent-tab, #parameters, #chat-settings, #lora, #training-tab, #model-tab, #session-tab, #character-tab {
#notebook-parent-tab, #parameters, #chat-settings, #lora, #training-tab, #model-tab, #session-tab, #character-tab, #image-ai-tab {
padding: 1rem;
}
@ -244,37 +244,46 @@ button {
font-size: 100% !important;
}
.pretty_scrollbar::-webkit-scrollbar {
.pretty_scrollbar::-webkit-scrollbar,
#image-history-gallery > :nth-child(2)::-webkit-scrollbar {
width: 8px;
height: 8px;
}
.pretty_scrollbar::-webkit-scrollbar-track {
.pretty_scrollbar::-webkit-scrollbar-track,
#image-history-gallery > :nth-child(2)::-webkit-scrollbar-track {
background: transparent;
}
.pretty_scrollbar::-webkit-scrollbar-thumb,
.pretty_scrollbar::-webkit-scrollbar-thumb:hover {
.pretty_scrollbar::-webkit-scrollbar-thumb:hover,
#image-history-gallery > :nth-child(2)::-webkit-scrollbar-thumb,
#image-history-gallery > :nth-child(2)::-webkit-scrollbar-thumb:hover {
background: var(--neutral-300);
border-radius: 30px;
}
.dark .pretty_scrollbar::-webkit-scrollbar-thumb,
.dark .pretty_scrollbar::-webkit-scrollbar-thumb:hover {
.dark .pretty_scrollbar::-webkit-scrollbar-thumb:hover,
.dark #image-history-gallery > :nth-child(2)::-webkit-scrollbar-thumb,
.dark #image-history-gallery > :nth-child(2)::-webkit-scrollbar-thumb:hover {
background: rgb(255 255 255 / 6.25%);
border-radius: 10px;
}
.pretty_scrollbar::-webkit-resizer {
.pretty_scrollbar::-webkit-resizer,
#image-history-gallery > :nth-child(2)::-webkit-resizer {
background: #c5c5d2;
}
.dark .pretty_scrollbar::-webkit-resizer {
.dark .pretty_scrollbar::-webkit-resizer,
.dark #image-history-gallery > :nth-child(2)::-webkit-resizer {
background: #ccc;
border-radius: 10px;
}
.pretty_scrollbar::-webkit-scrollbar-corner {
.pretty_scrollbar::-webkit-scrollbar-corner,
#image-history-gallery > :nth-child(2)::-webkit-scrollbar-corner {
background: transparent;
}
@ -1674,3 +1683,72 @@ button:focus {
.dark .sidebar-vertical-separator {
border-bottom: 1px solid rgb(255 255 255 / 10%);
}
button#swap-height-width {
position: absolute;
top: -50px;
right: 0;
border: 0;
}
#image-output-gallery, #image-output-gallery > :nth-child(2) {
height: calc(100vh - 83px);
max-height: calc(100vh - 83px);
}
#image-history-gallery, #image-history-gallery > :nth-child(2) {
height: calc(100vh - 174px);
max-height: calc(100vh - 174px);
}
/* Additional CSS for the paginated image gallery */
/* Page info styling */
#image-page-info {
display: flex;
align-items: center;
justify-content: center;
min-width: 200px;
font-size: 0.9em;
color: var(--body-text-color-subdued);
}
/* Settings display panel */
#image-ai-tab .settings-display-panel {
background: var(--background-fill-secondary);
padding: 12px;
border-radius: 8px;
font-size: 0.9em;
max-height: 300px;
overflow-y: auto;
margin-top: 8px;
}
/* Gallery status message */
#image-ai-tab .gallery-status {
color: var(--color-accent);
font-size: 0.85em;
margin-top: 4px;
}
/* Pagination button row alignment */
#image-ai-tab .pagination-controls {
display: flex;
align-items: center;
gap: 8px;
flex-wrap: wrap;
}
/* Selected image preview container */
#image-ai-tab .selected-preview-container {
border: 1px solid var(--border-color-primary);
border-radius: 8px;
padding: 8px;
background: var(--background-fill-secondary);
}
/* Fix a gr.Markdown UI glitch when clicking Next in the
* Image AI > Gallery tab */
.min.svelte-1yrv54 {
min-height: 0;
}

View file

@ -0,0 +1,20 @@
# Image Generation Tutorial
This feature allows you to generate images using high-speed models like Z-Image-Turbo directly within the web UI.
## How to use
1. Click on the **Image AI** tab at the top of the interface.
2. Select the **Model** sub-tab.
3. Copy and paste the following link into the **Download model** box:
```
https://huggingface.co/Tongyi-MAI/Z-Image-Turbo
```
4. Click the **Download** button and wait for the confirmation message.
5. In the **Model** dropdown menu, select the model you just downloaded (if you don't see it, click the 🔄 refresh button).
6. Click **Load**.
7. Go to the **Generate** sub-tab, type a prompt, and click **GENERATE**.
> **Note for Z-Image-Turbo:** For the best results with this specific model, keep the **CFG Scale** slider at **0**.

159
modules/image_models.py Normal file
View file

@ -0,0 +1,159 @@
import time
import modules.shared as shared
from modules.logging_colors import logger
from modules.torch_utils import get_device
from modules.utils import resolve_model_path
def get_quantization_config(quant_method):
"""
Get the appropriate quantization config based on the selected method.
Args:
quant_method: One of 'none', 'bnb-8bit', 'bnb-4bit'
Returns:
PipelineQuantizationConfig or None
"""
import torch
from diffusers import BitsAndBytesConfig, QuantoConfig
from diffusers.quantizers import PipelineQuantizationConfig
if quant_method == 'none' or not quant_method:
return None
# Bitsandbytes 8-bit quantization
elif quant_method == 'bnb-8bit':
return PipelineQuantizationConfig(
quant_mapping={
"transformer": BitsAndBytesConfig(
load_in_8bit=True
)
}
)
# Bitsandbytes 4-bit quantization
elif quant_method == 'bnb-4bit':
return PipelineQuantizationConfig(
quant_mapping={
"transformer": BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True
)
}
)
else:
logger.warning(f"Unknown quantization method: {quant_method}. Loading without quantization.")
return None
def get_pipeline_type(pipe):
"""
Detect the pipeline type based on the loaded pipeline class.
Returns:
str: 'zimage', 'qwenimage', or 'unknown'
"""
class_name = pipe.__class__.__name__
if class_name == 'ZImagePipeline':
return 'zimage'
elif class_name == 'QwenImagePipeline':
return 'qwenimage'
else:
return 'unknown'
def load_image_model(model_name, dtype='bfloat16', attn_backend='sdpa', cpu_offload=False, compile_model=False, quant_method='none'):
"""
Load a diffusers image generation model.
Args:
model_name: Name of the model directory
dtype: 'bfloat16' or 'float16'
attn_backend: 'sdpa', 'flash_attention_2', or 'flash_attention_3'
cpu_offload: Enable CPU offloading for low VRAM
compile_model: Compile the model for faster inference (slow first run)
quant_method: Quantization method - 'none', 'bnb-8bit', 'bnb-4bit'
"""
import torch
from diffusers import DiffusionPipeline
logger.info(f"Loading image model \"{model_name}\" with quantization: {quant_method}")
t0 = time.time()
dtype_map = {"bfloat16": torch.bfloat16, "float16": torch.float16}
target_dtype = dtype_map.get(dtype, torch.bfloat16)
model_path = resolve_model_path(model_name, image_model=True)
try:
# Get quantization config based on selected method
pipeline_quant_config = get_quantization_config(quant_method)
# Load the pipeline
load_kwargs = {
"torch_dtype": target_dtype,
"low_cpu_mem_usage": True,
}
if pipeline_quant_config is not None:
load_kwargs["quantization_config"] = pipeline_quant_config
# Use DiffusionPipeline for automatic pipeline detection
# This handles both ZImagePipeline and QwenImagePipeline
pipe = DiffusionPipeline.from_pretrained(
str(model_path),
**load_kwargs
)
pipeline_type = get_pipeline_type(pipe)
if not cpu_offload:
pipe.to(get_device())
# Set attention backend (if supported by the pipeline)
if hasattr(pipe, 'transformer') and hasattr(pipe.transformer, 'set_attention_backend'):
if attn_backend == 'flash_attention_2':
pipe.transformer.set_attention_backend("flash")
elif attn_backend == 'flash_attention_3':
pipe.transformer.set_attention_backend("_flash_3")
# sdpa is the default, no action needed
if compile_model:
if hasattr(pipe, 'transformer') and hasattr(pipe.transformer, 'compile'):
logger.info("Compiling model (first run will be slow)...")
pipe.transformer.compile()
if cpu_offload:
pipe.enable_model_cpu_offload()
shared.image_model = pipe
shared.image_model_name = model_name
shared.image_pipeline_type = pipeline_type
logger.info(f"Loaded image model \"{model_name}\" in {(time.time() - t0):.2f} seconds.")
return pipe
except Exception as e:
logger.error(f"Failed to load image model: {str(e)}")
return None
def unload_image_model():
"""Unload the current image model and free VRAM."""
if shared.image_model is None:
return
del shared.image_model
shared.image_model = None
shared.image_model_name = 'None'
shared.image_pipeline_type = None
from modules.torch_utils import clear_torch_cache
clear_torch_cache()
logger.info("Image model unloaded.")

View file

@ -11,7 +11,7 @@ import yaml
from modules.logging_colors import logger
from modules.presets import default_preset
# Model variables
# Text model variables
model = None
tokenizer = None
model_name = 'None'
@ -20,6 +20,11 @@ is_multimodal = False
model_dirty_from_training = False
lora_names = []
# Image model variables
image_model = None
image_model_name = 'None'
image_pipeline_type = None
# Generation variables
stop_everything = False
generation_lock = None
@ -46,6 +51,18 @@ group.add_argument('--extensions', type=str, nargs='+', help='The list of extens
group.add_argument('--verbose', action='store_true', help='Print the prompts to the terminal.')
group.add_argument('--idle-timeout', type=int, default=0, help='Unload model after this many minutes of inactivity. It will be automatically reloaded when you try to use it again.')
# Image generation
group = parser.add_argument_group('Image model')
group.add_argument('--image-model', type=str, help='Name of the image model to select on startup (overrides saved setting).')
group.add_argument('--image-model-dir', type=str, default='user_data/image_models', help='Path to directory with all the image models.')
group.add_argument('--image-dtype', type=str, default=None, choices=['bfloat16', 'float16'], help='Data type for image model.')
group.add_argument('--image-attn-backend', type=str, default=None, choices=['sdpa', 'flash_attention_2', 'flash_attention_3'], help='Attention backend for image model.')
group.add_argument('--image-cpu-offload', action='store_true', help='Enable CPU offloading for image model.')
group.add_argument('--image-compile', action='store_true', help='Compile the image model for faster inference.')
group.add_argument('--image-quant', type=str, default=None,
choices=['none', 'bnb-8bit', 'bnb-4bit'],
help='Quantization method for image model.')
# Model loader
group = parser.add_argument_group('Model loader')
group.add_argument('--loader', type=str, help='Choose the model loader manually, otherwise, it will get autodetected. Valid options: Transformers, llama.cpp, ExLlamav3_HF, ExLlamav2_HF, ExLlamav2, TensorRT-LLM.')
@ -290,6 +307,24 @@ settings = {
# Extensions
'default_extensions': [],
# Image generation settings
'image_prompt': '',
'image_neg_prompt': '',
'image_width': 1024,
'image_height': 1024,
'image_aspect_ratio': '1:1 Square',
'image_steps': 9,
'image_cfg_scale': 0.0,
'image_seed': -1,
'image_batch_size': 1,
'image_batch_count': 1,
'image_model_menu': 'None',
'image_dtype': 'bfloat16',
'image_attn_backend': 'sdpa',
'image_cpu_offload': False,
'image_compile': False,
'image_quant': 'none',
}
default_settings = copy.deepcopy(settings)
@ -314,6 +349,22 @@ def do_cmd_flags_warnings():
logger.warning('\nThe multi-user mode is highly experimental and should not be shared publicly.')
def apply_image_model_cli_overrides():
"""Apply command-line overrides for image model settings."""
if args.image_model is not None:
settings['image_model_menu'] = args.image_model
if args.image_dtype is not None:
settings['image_dtype'] = args.image_dtype
if args.image_attn_backend is not None:
settings['image_attn_backend'] = args.image_attn_backend
if args.image_cpu_offload:
settings['image_cpu_offload'] = True
if args.image_compile:
settings['image_compile'] = True
if args.image_quant is not None:
settings['image_quant'] = args.image_quant
def fix_loader_name(name):
if not name:
return name

View file

@ -280,6 +280,26 @@ def list_interface_input_elements():
'include_past_attachments',
]
# Image generation elements
elements += [
'image_prompt',
'image_neg_prompt',
'image_width',
'image_height',
'image_aspect_ratio',
'image_steps',
'image_cfg_scale',
'image_seed',
'image_batch_size',
'image_batch_count',
'image_model_menu',
'image_dtype',
'image_attn_backend',
'image_compile',
'image_cpu_offload',
'image_quant',
]
return elements
@ -509,7 +529,25 @@ def setup_auto_save():
'theme_state',
'show_two_notebook_columns',
'paste_to_attachment',
'include_past_attachments'
'include_past_attachments',
# Image generation tab (ui_image_generation.py)
'image_prompt',
'image_neg_prompt',
'image_width',
'image_height',
'image_aspect_ratio',
'image_steps',
'image_cfg_scale',
'image_seed',
'image_batch_size',
'image_batch_count',
'image_model_menu',
'image_dtype',
'image_attn_backend',
'image_compile',
'image_cpu_offload',
'image_quant',
]
for element_name in change_elements:

View file

@ -0,0 +1,847 @@
import json
import os
import time
import traceback
from datetime import datetime
from pathlib import Path
import gradio as gr
import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from modules import shared, ui, utils
from modules.image_models import (
get_pipeline_type,
load_image_model,
unload_image_model
)
from modules.logging_colors import logger
from modules.text_generation import stop_everything_event
from modules.torch_utils import get_device
from modules.utils import gradio
ASPECT_RATIOS = {
"1:1 Square": (1, 1),
"16:9 Cinema": (16, 9),
"9:16 Mobile": (9, 16),
"4:3 Photo": (4, 3),
"Custom": None,
}
STEP = 16
IMAGES_PER_PAGE = 64
# Settings keys to save in PNG metadata (Generate tab only)
METADATA_SETTINGS_KEYS = [
'image_prompt',
'image_neg_prompt',
'image_width',
'image_height',
'image_aspect_ratio',
'image_steps',
'image_seed',
'image_batch_size',
'image_batch_count',
'image_cfg_scale',
]
# Cache for all image paths
_image_cache = []
_cache_timestamp = 0
def round_to_step(value, step=STEP):
return round(value / step) * step
def clamp(value, min_val, max_val):
return max(min_val, min(max_val, value))
def apply_aspect_ratio(aspect_ratio, current_width, current_height):
if aspect_ratio == "Custom" or aspect_ratio not in ASPECT_RATIOS:
return current_width, current_height
w_ratio, h_ratio = ASPECT_RATIOS[aspect_ratio]
if w_ratio == h_ratio:
base = min(current_width, current_height)
new_width = base
new_height = base
elif w_ratio < h_ratio:
new_width = current_width
new_height = round_to_step(current_width * h_ratio / w_ratio)
else:
new_height = current_height
new_width = round_to_step(current_height * w_ratio / h_ratio)
new_width = clamp(new_width, 256, 2048)
new_height = clamp(new_height, 256, 2048)
return int(new_width), int(new_height)
def update_height_from_width(width, aspect_ratio):
if aspect_ratio == "Custom" or aspect_ratio not in ASPECT_RATIOS:
return gr.update()
w_ratio, h_ratio = ASPECT_RATIOS[aspect_ratio]
new_height = round_to_step(width * h_ratio / w_ratio)
new_height = clamp(new_height, 256, 2048)
return int(new_height)
def update_width_from_height(height, aspect_ratio):
if aspect_ratio == "Custom" or aspect_ratio not in ASPECT_RATIOS:
return gr.update()
w_ratio, h_ratio = ASPECT_RATIOS[aspect_ratio]
new_width = round_to_step(height * w_ratio / h_ratio)
new_width = clamp(new_width, 256, 2048)
return int(new_width)
def swap_dimensions_and_update_ratio(width, height, aspect_ratio):
new_width, new_height = height, width
new_ratio = "Custom"
for name, ratios in ASPECT_RATIOS.items():
if ratios is None:
continue
w_r, h_r = ratios
expected_height = new_width * h_r / w_r
if abs(expected_height - new_height) < STEP:
new_ratio = name
break
return new_width, new_height, new_ratio
def build_generation_metadata(state, actual_seed):
"""Build metadata dict from generation settings."""
metadata = {}
for key in METADATA_SETTINGS_KEYS:
if key in state:
metadata[key] = state[key]
# Store the actual seed used (not -1)
metadata['image_seed'] = actual_seed
metadata['generated_at'] = datetime.now().isoformat()
metadata['model'] = shared.image_model_name
return metadata
def save_generated_images(images, state, actual_seed):
"""Save images with generation metadata embedded in PNG."""
date_str = datetime.now().strftime("%Y-%m-%d")
folder_path = os.path.join("user_data", "image_outputs", date_str)
os.makedirs(folder_path, exist_ok=True)
metadata = build_generation_metadata(state, actual_seed)
metadata_json = json.dumps(metadata, ensure_ascii=False)
for idx, img in enumerate(images):
timestamp = datetime.now().strftime("%H-%M-%S")
filename = f"{timestamp}_{actual_seed:010d}_{idx:03d}.png"
filepath = os.path.join(folder_path, filename)
# Create PNG metadata
png_info = PngInfo()
png_info.add_text("image_gen_settings", metadata_json)
# Save with metadata
img.save(filepath, pnginfo=png_info)
def read_image_metadata(image_path):
"""Read generation metadata from PNG file."""
try:
with Image.open(image_path) as img:
if hasattr(img, 'text') and 'image_gen_settings' in img.text:
return json.loads(img.text['image_gen_settings'])
except Exception as e:
logger.debug(f"Could not read metadata from {image_path}: {e}")
return None
def format_metadata_for_display(metadata):
"""Format metadata as readable text."""
if not metadata:
return "No generation settings found in this image."
lines = ["**Generation Settings**", ""]
# Display in a nice order
display_order = [
('image_prompt', 'Prompt'),
('image_neg_prompt', 'Negative Prompt'),
('image_width', 'Width'),
('image_height', 'Height'),
('image_aspect_ratio', 'Aspect Ratio'),
('image_steps', 'Steps'),
('image_cfg_scale', 'CFG Scale'),
('image_seed', 'Seed'),
('image_batch_size', 'Batch Size'),
('image_batch_count', 'Batch Count'),
('model', 'Model'),
('generated_at', 'Generated At'),
]
for key, label in display_order:
if key in metadata:
value = metadata[key]
if key in ['image_prompt', 'image_neg_prompt'] and value:
# Truncate long prompts for display
if len(str(value)) > 200:
value = str(value)[:200] + "..."
lines.append(f"**{label}:** {value}")
return "\n\n".join(lines)
def get_all_history_images(force_refresh=False):
"""Get all history images sorted by modification time (newest first). Uses caching."""
global _image_cache, _cache_timestamp
output_dir = os.path.join("user_data", "image_outputs")
if not os.path.exists(output_dir):
return []
# Check if we need to refresh cache
current_time = time.time()
if not force_refresh and _image_cache and (current_time - _cache_timestamp) < 2:
return _image_cache
image_files = []
for root, _, files in os.walk(output_dir):
for file in files:
if file.endswith((".png", ".jpg", ".jpeg")):
full_path = os.path.join(root, file)
image_files.append((full_path, os.path.getmtime(full_path)))
image_files.sort(key=lambda x: x[1], reverse=True)
_image_cache = [x[0] for x in image_files]
_cache_timestamp = current_time
return _image_cache
def get_paginated_images(page=0, force_refresh=False):
"""Get images for a specific page."""
all_images = get_all_history_images(force_refresh)
total_images = len(all_images)
total_pages = max(1, (total_images + IMAGES_PER_PAGE - 1) // IMAGES_PER_PAGE)
# Clamp page to valid range
page = max(0, min(page, total_pages - 1))
start_idx = page * IMAGES_PER_PAGE
end_idx = min(start_idx + IMAGES_PER_PAGE, total_images)
page_images = all_images[start_idx:end_idx]
return page_images, page, total_pages, total_images
def get_initial_page_info():
"""Get page info string for initial load."""
_, page, total_pages, total_images = get_paginated_images(0)
return f"Page {page + 1} of {total_pages} ({total_images} total images)"
def refresh_gallery(current_page=0):
"""Refresh gallery with current page."""
images, page, total_pages, total_images = get_paginated_images(current_page, force_refresh=True)
page_info = f"Page {page + 1} of {total_pages} ({total_images} total images)"
return images, page, page_info
def go_to_page(page_num, current_page):
"""Go to a specific page (1-indexed input)."""
try:
page = int(page_num) - 1 # Convert to 0-indexed
except (ValueError, TypeError):
page = current_page
images, page, total_pages, total_images = get_paginated_images(page)
page_info = f"Page {page + 1} of {total_pages} ({total_images} total images)"
return images, page, page_info
def next_page(current_page):
"""Go to next page."""
images, page, total_pages, total_images = get_paginated_images(current_page + 1)
page_info = f"Page {page + 1} of {total_pages} ({total_images} total images)"
return images, page, page_info
def prev_page(current_page):
"""Go to previous page."""
images, page, total_pages, total_images = get_paginated_images(current_page - 1)
page_info = f"Page {page + 1} of {total_pages} ({total_images} total images)"
return images, page, page_info
def on_gallery_select(evt: gr.SelectData, current_page):
"""Handle image selection from gallery."""
if evt.index is None:
return "", "Select an image to view its settings"
# Get the current page's images to find the actual file path
all_images = get_all_history_images()
total_images = len(all_images)
# Calculate the actual index in the full list
start_idx = current_page * IMAGES_PER_PAGE
actual_idx = start_idx + evt.index
if actual_idx >= total_images:
return "", "Image not found"
image_path = all_images[actual_idx]
metadata = read_image_metadata(image_path)
metadata_display = format_metadata_for_display(metadata)
return image_path, metadata_display
def send_to_generate(selected_image_path):
"""Load settings from selected image and return updates for all Generate tab inputs."""
if not selected_image_path or not os.path.exists(selected_image_path):
return [gr.update()] * 10 + ["No image selected"]
metadata = read_image_metadata(selected_image_path)
if not metadata:
return [gr.update()] * 10 + ["No settings found in this image"]
# Return updates for each input element in order
updates = [
gr.update(value=metadata.get('image_prompt', '')),
gr.update(value=metadata.get('image_neg_prompt', '')),
gr.update(value=metadata.get('image_width', 1024)),
gr.update(value=metadata.get('image_height', 1024)),
gr.update(value=metadata.get('image_aspect_ratio', '1:1 Square')),
gr.update(value=metadata.get('image_steps', 9)),
gr.update(value=metadata.get('image_seed', -1)),
gr.update(value=metadata.get('image_batch_size', 1)),
gr.update(value=metadata.get('image_batch_count', 1)),
gr.update(value=metadata.get('image_cfg_scale', 0.0)),
]
status = f"✓ Settings loaded from image (seed: {metadata.get('image_seed', 'unknown')})"
return updates + [status]
def read_dropped_image_metadata(image_path):
"""Read metadata from a dropped/uploaded image."""
if not image_path:
return "Drop an image to view its generation settings."
metadata = read_image_metadata(image_path)
return format_metadata_for_display(metadata)
def create_ui():
if shared.settings['image_model_menu'] != 'None':
shared.image_model_name = shared.settings['image_model_menu']
with gr.Tab("Image AI", elem_id="image-ai-tab"):
with gr.Tabs():
# TAB 1: GENERATE
with gr.TabItem("Generate"):
with gr.Row():
with gr.Column(scale=4, min_width=350):
shared.gradio['image_prompt'] = gr.Textbox(
label="Prompt",
placeholder="Describe your imagination...",
lines=3,
autofocus=True,
value=shared.settings['image_prompt']
)
shared.gradio['image_neg_prompt'] = gr.Textbox(
label="Negative Prompt",
placeholder="Low quality...",
lines=3,
value=shared.settings['image_neg_prompt']
)
shared.gradio['image_generate_btn'] = gr.Button("Generate", variant="primary", size="lg")
shared.gradio['image_stop_btn'] = gr.Button("Stop", size="lg", visible=False)
gr.HTML("<hr style='border-top: 1px solid #444; margin: 20px 0;'>")
gr.Markdown("### Dimensions")
with gr.Row():
with gr.Column():
shared.gradio['image_width'] = gr.Slider(256, 2048, value=shared.settings['image_width'], step=STEP, label="Width")
with gr.Column():
shared.gradio['image_height'] = gr.Slider(256, 2048, value=shared.settings['image_height'], step=STEP, label="Height")
shared.gradio['image_swap_btn'] = gr.Button("⇄ Swap", elem_classes='refresh-button', scale=0, min_width=80, elem_id="swap-height-width")
with gr.Row():
shared.gradio['image_aspect_ratio'] = gr.Radio(
choices=["1:1 Square", "16:9 Cinema", "9:16 Mobile", "4:3 Photo", "Custom"],
value=shared.settings['image_aspect_ratio'],
label="Aspect Ratio",
interactive=True
)
gr.Markdown("### Config")
with gr.Row():
with gr.Column():
shared.gradio['image_steps'] = gr.Slider(1, 100, value=shared.settings['image_steps'], step=1, label="Steps")
shared.gradio['image_cfg_scale'] = gr.Slider(
0.0, 10.0,
value=shared.settings['image_cfg_scale'],
step=0.1,
label="CFG Scale",
info="Z-Image Turbo: 0.0 | Qwen: 4.0"
)
shared.gradio['image_seed'] = gr.Number(label="Seed", value=shared.settings['image_seed'], precision=0, info="-1 = Random")
with gr.Column():
shared.gradio['image_batch_size'] = gr.Slider(1, 32, value=shared.settings['image_batch_size'], step=1, label="Batch Size (VRAM Heavy)", info="Generates N images at once.")
shared.gradio['image_batch_count'] = gr.Slider(1, 128, value=shared.settings['image_batch_count'], step=1, label="Sequential Count (Loop)", info="Repeats the generation N times.")
with gr.Column(scale=6, min_width=500):
with gr.Column(elem_classes=["viewport-container"]):
shared.gradio['image_output_gallery'] = gr.Gallery(label="Output", show_label=False, columns=2, rows=2, height="80vh", object_fit="contain", preview=True, elem_id="image-output-gallery")
# TAB 2: GALLERY (with pagination)
with gr.TabItem("Gallery"):
with gr.Row():
with gr.Column(scale=3):
# Pagination controls
with gr.Row():
shared.gradio['image_refresh_history'] = gr.Button("🔄 Refresh", elem_classes="refresh-button")
shared.gradio['image_prev_page'] = gr.Button("◀ Prev", elem_classes="refresh-button")
shared.gradio['image_page_info'] = gr.Markdown(value=get_initial_page_info, elem_id="image-page-info")
shared.gradio['image_next_page'] = gr.Button("Next ▶", elem_classes="refresh-button")
shared.gradio['image_page_input'] = gr.Number(value=1, label="Page", precision=0, minimum=1, scale=0, min_width=80)
shared.gradio['image_go_to_page'] = gr.Button("Go", elem_classes="refresh-button", scale=0, min_width=50)
# State for current page and selected image path
shared.gradio['image_current_page'] = gr.State(value=0)
shared.gradio['image_selected_path'] = gr.State(value="")
# Paginated gallery using gr.Gallery
shared.gradio['image_history_gallery'] = gr.Gallery(
value=lambda: get_paginated_images(0)[0],
label="Image History",
show_label=False,
columns=6,
object_fit="cover",
height="auto",
allow_preview=True,
elem_id="image-history-gallery"
)
with gr.Column(scale=1):
gr.Markdown("### Selected Image")
shared.gradio['image_settings_display'] = gr.Markdown("Select an image to view its settings")
shared.gradio['image_send_to_generate'] = gr.Button("Send to Generate", variant="primary")
shared.gradio['image_gallery_status'] = gr.Markdown("")
gr.Markdown("### Import Image")
shared.gradio['image_drop_upload'] = gr.Image(
label="Drop image here to view settings",
type="filepath",
height=150
)
# TAB 3: MODEL
with gr.TabItem("Model"):
with gr.Row():
with gr.Column():
with gr.Row():
shared.gradio['image_model_menu'] = gr.Dropdown(
choices=utils.get_available_image_models(),
value=shared.settings['image_model_menu'],
label='Model',
elem_classes='slim-dropdown'
)
shared.gradio['image_refresh_models'] = gr.Button("🔄", elem_classes='refresh-button', scale=0, min_width=40)
shared.gradio['image_load_model'] = gr.Button("Load", variant='primary', elem_classes='refresh-button')
shared.gradio['image_unload_model'] = gr.Button("Unload", elem_classes='refresh-button')
gr.Markdown("## Settings")
with gr.Row():
with gr.Column():
shared.gradio['image_quant'] = gr.Dropdown(
label='Quantization',
choices=['none', 'bnb-8bit', 'bnb-4bit'],
value=shared.settings['image_quant'],
info='Quantization method for reduced VRAM usage. Quanto supports lower precisions (2-bit, 4-bit, 8-bit).'
)
shared.gradio['image_dtype'] = gr.Dropdown(
choices=['bfloat16', 'float16'],
value=shared.settings['image_dtype'],
label='Data Type',
info='bfloat16 recommended for modern GPUs'
)
shared.gradio['image_attn_backend'] = gr.Dropdown(
choices=['sdpa', 'flash_attention_2', 'flash_attention_3'],
value=shared.settings['image_attn_backend'],
label='Attention Backend',
info='SDPA is default. Flash Attention requires compatible GPU.'
)
with gr.Column():
shared.gradio['image_compile'] = gr.Checkbox(
value=shared.settings['image_compile'],
label='Compile Model',
info='Faster inference after first run. First run will be slow.'
)
shared.gradio['image_cpu_offload'] = gr.Checkbox(
value=shared.settings['image_cpu_offload'],
label='CPU Offload',
info='Enable for low VRAM GPUs. Slower but uses less memory.'
)
with gr.Column():
shared.gradio['image_download_path'] = gr.Textbox(
label="Download model",
placeholder="Tongyi-MAI/Z-Image-Turbo",
info="Enter HuggingFace path. Use : for branch, e.g. user/model:main"
)
shared.gradio['image_download_btn'] = gr.Button("Download", variant='primary')
shared.gradio['image_model_status'] = gr.Markdown(
value=f"Model: **{shared.settings['image_model_menu']}** (not loaded)" if shared.settings['image_model_menu'] != 'None' else "No model selected"
)
def create_event_handlers():
# Dimension controls
shared.gradio['image_aspect_ratio'].change(
apply_aspect_ratio,
gradio('image_aspect_ratio', 'image_width', 'image_height'),
gradio('image_width', 'image_height'),
show_progress=False
)
shared.gradio['image_width'].release(
update_height_from_width,
gradio('image_width', 'image_aspect_ratio'),
gradio('image_height'),
show_progress=False
)
shared.gradio['image_height'].release(
update_width_from_height,
gradio('image_height', 'image_aspect_ratio'),
gradio('image_width'),
show_progress=False
)
shared.gradio['image_swap_btn'].click(
swap_dimensions_and_update_ratio,
gradio('image_width', 'image_height', 'image_aspect_ratio'),
gradio('image_width', 'image_height', 'image_aspect_ratio'),
show_progress=False
)
# Generation
shared.gradio['image_generate_btn'].click(
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
lambda: [gr.update(visible=True), gr.update(visible=False)], None, gradio('image_stop_btn', 'image_generate_btn')).then(
generate, gradio('interface_state'), gradio('image_output_gallery'), show_progress=False).then(
lambda: [gr.update(visible=False), gr.update(visible=True)], None, gradio('image_stop_btn', 'image_generate_btn'))
shared.gradio['image_prompt'].submit(
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
lambda: [gr.update(visible=True), gr.update(visible=False)], None, gradio('image_stop_btn', 'image_generate_btn')).then(
generate, gradio('interface_state'), gradio('image_output_gallery'), show_progress=False).then(
lambda: [gr.update(visible=False), gr.update(visible=True)], None, gradio('image_stop_btn', 'image_generate_btn'))
shared.gradio['image_neg_prompt'].submit(
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
lambda: [gr.update(visible=True), gr.update(visible=False)], None, gradio('image_stop_btn', 'image_generate_btn')).then(
generate, gradio('interface_state'), gradio('image_output_gallery'), show_progress=False).then(
lambda: [gr.update(visible=False), gr.update(visible=True)], None, gradio('image_stop_btn', 'image_generate_btn'))
# Stop button
shared.gradio['image_stop_btn'].click(
stop_everything_event, None, None, show_progress=False
)
# Model management
shared.gradio['image_refresh_models'].click(
lambda: gr.update(choices=utils.get_available_image_models()),
None,
gradio('image_model_menu'),
show_progress=False
)
shared.gradio['image_load_model'].click(
load_image_model_wrapper,
gradio('image_model_menu', 'image_dtype', 'image_attn_backend', 'image_cpu_offload', 'image_compile', 'image_quant'),
gradio('image_model_status'),
show_progress=True
)
shared.gradio['image_unload_model'].click(
unload_image_model_wrapper,
None,
gradio('image_model_status'),
show_progress=False
)
shared.gradio['image_download_btn'].click(
download_image_model_wrapper,
gradio('image_download_path'),
gradio('image_model_status', 'image_model_menu'),
show_progress=True
)
# Gallery pagination handlers
shared.gradio['image_refresh_history'].click(
refresh_gallery,
gradio('image_current_page'),
gradio('image_history_gallery', 'image_current_page', 'image_page_info'),
show_progress=False
)
shared.gradio['image_next_page'].click(
next_page,
gradio('image_current_page'),
gradio('image_history_gallery', 'image_current_page', 'image_page_info'),
show_progress=False
)
shared.gradio['image_prev_page'].click(
prev_page,
gradio('image_current_page'),
gradio('image_history_gallery', 'image_current_page', 'image_page_info'),
show_progress=False
)
shared.gradio['image_go_to_page'].click(
go_to_page,
gradio('image_page_input', 'image_current_page'),
gradio('image_history_gallery', 'image_current_page', 'image_page_info'),
show_progress=False
)
# Image selection from gallery
shared.gradio['image_history_gallery'].select(
on_gallery_select,
gradio('image_current_page'),
gradio('image_selected_path', 'image_settings_display'),
show_progress=False
)
# Send to Generate
shared.gradio['image_send_to_generate'].click(
send_to_generate,
gradio('image_selected_path'),
gradio(
'image_prompt',
'image_neg_prompt',
'image_width',
'image_height',
'image_aspect_ratio',
'image_steps',
'image_seed',
'image_batch_size',
'image_batch_count',
'image_cfg_scale',
'image_gallery_status'
),
show_progress=False
)
shared.gradio['image_drop_upload'].change(
read_dropped_image_metadata,
gradio('image_drop_upload'),
gradio('image_settings_display'),
show_progress=False
)
def generate(state):
"""
Generate images using the loaded model.
Automatically adjusts parameters based on pipeline type.
"""
import torch
from modules.torch_utils import clear_torch_cache
clear_torch_cache()
try:
model_name = state['image_model_menu']
if not model_name or model_name == 'None':
logger.error("No image model selected. Go to the Model tab and select a model.")
return []
if shared.image_model is None:
result = load_image_model(
model_name,
dtype=state['image_dtype'],
attn_backend=state['image_attn_backend'],
cpu_offload=state['image_cpu_offload'],
compile_model=state['image_compile'],
quant_method=state['image_quant']
)
if result is None:
logger.error(f"Failed to load model `{model_name}`.")
return []
shared.image_model_name = model_name
seed = state['image_seed']
if seed == -1:
seed = np.random.randint(0, 2**32 - 1)
device = get_device()
if device is None:
device = "cpu"
generator = torch.Generator(device).manual_seed(int(seed))
all_images = []
# Get pipeline type for parameter adjustment
pipeline_type = getattr(shared, 'image_pipeline_type', None)
if pipeline_type is None:
pipeline_type = get_pipeline_type(shared.image_model)
# Process Prompt
prompt = state['image_prompt']
# Apply "Positive Magic" for Qwen models only
if pipeline_type == 'qwenimage':
magic_suffix = ", Ultra HD, 4K, cinematic composition"
# Avoid duplication if user already added it
if magic_suffix.strip(", ") not in prompt:
prompt += magic_suffix
# Reset stop flag at start
shared.stop_everything = False
# Callback to check for interruption during diffusion steps
def interrupt_callback(pipe, step_index, timestep, callback_kwargs):
if shared.stop_everything:
pipe._interrupt = True
return callback_kwargs
# Build generation kwargs
gen_kwargs = {
"prompt": prompt,
"negative_prompt": state['image_neg_prompt'],
"height": int(state['image_height']),
"width": int(state['image_width']),
"num_inference_steps": int(state['image_steps']),
"num_images_per_prompt": int(state['image_batch_size']),
"generator": generator,
"callback_on_step_end": interrupt_callback,
}
# Add pipeline-specific parameters for CFG
cfg_val = state.get('image_cfg_scale', 0.0)
if pipeline_type == 'qwenimage':
# Qwen-Image uses true_cfg_scale (typically 4.0)
gen_kwargs["true_cfg_scale"] = cfg_val
else:
# Z-Image and others use guidance_scale (typically 0.0 for Turbo)
gen_kwargs["guidance_scale"] = cfg_val
t0 = time.time()
for i in range(int(state['image_batch_count'])):
if shared.stop_everything:
break
generator.manual_seed(int(seed + i))
batch_results = shared.image_model(**gen_kwargs).images
all_images.extend(batch_results)
t1 = time.time()
save_generated_images(all_images, state, seed)
total_images = int(state['image_batch_count']) * int(state['image_batch_size'])
total_steps = state["image_steps"] * int(state['image_batch_count'])
logger.info(f'Generated {total_images} {"image" if total_images == 1 else "images"} in {(t1 - t0):.2f} seconds ({total_steps / (t1 - t0):.2f} steps/s, seed {seed})')
return all_images
except Exception as e:
logger.error(f"Image generation failed: {e}")
traceback.print_exc()
return []
def load_image_model_wrapper(model_name, dtype, attn_backend, cpu_offload, compile_model, quant_method):
if not model_name or model_name == 'None':
yield "No model selected"
return
try:
yield f"Loading `{model_name}`..."
unload_image_model()
result = load_image_model(
model_name,
dtype=dtype,
attn_backend=attn_backend,
cpu_offload=cpu_offload,
compile_model=compile_model,
quant_method=quant_method
)
if result is not None:
shared.image_model_name = model_name
yield f"✓ Loaded **{model_name}** (quantization: {quant_method})"
else:
yield f"✗ Failed to load `{model_name}`"
except Exception:
yield f"Error:\n```\n{traceback.format_exc()}\n```"
def unload_image_model_wrapper():
previous_name = shared.image_model_name
unload_image_model()
if previous_name != 'None':
return f"Model: **{previous_name}** (unloaded)"
return "No model loaded"
def download_image_model_wrapper(model_path):
from huggingface_hub import snapshot_download
if not model_path:
yield "No model specified", gr.update()
return
try:
model_path = model_path.strip()
if model_path.startswith('https://huggingface.co/'):
model_path = model_path[len('https://huggingface.co/'):]
elif model_path.startswith('huggingface.co/'):
model_path = model_path[len('huggingface.co/'):]
if ':' in model_path:
model_id, branch = model_path.rsplit(':', 1)
else:
model_id, branch = model_path, 'main'
folder_name = model_id.replace('/', '_')
output_folder = Path(shared.args.image_model_dir) / folder_name
yield f"Downloading `{model_id}` (branch: {branch})...", gr.update()
snapshot_download(
repo_id=model_id,
revision=branch,
local_dir=output_folder,
local_dir_use_symlinks=False,
)
new_choices = utils.get_available_image_models()
yield f"✓ Downloaded to `{output_folder}`", gr.update(choices=new_choices, value=folder_name)
except Exception:
yield f"Error:\n```\n{traceback.format_exc()}\n```", gr.update()

View file

@ -86,7 +86,7 @@ def check_model_loaded():
return True, None
def resolve_model_path(model_name_or_path):
def resolve_model_path(model_name_or_path, image_model=False):
"""
Resolves a model path, checking for a direct path
before the default models directory.
@ -95,6 +95,8 @@ def resolve_model_path(model_name_or_path):
path_candidate = Path(model_name_or_path)
if path_candidate.exists():
return path_candidate
elif image_model:
return Path(f'{shared.args.image_model_dir}/{model_name_or_path}')
else:
return Path(f'{shared.args.model_dir}/{model_name_or_path}')
@ -153,6 +155,24 @@ def get_available_models():
return filtered_gguf_files + model_dirs
def get_available_image_models():
model_dir = Path(shared.args.image_model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
# Find valid model directories
model_dirs = []
for item in os.listdir(model_dir):
item_path = model_dir / item
if not item_path.is_dir():
continue
model_dirs.append(item)
model_dirs = sorted(model_dirs, key=natural_keys)
return model_dirs
def get_available_ggufs():
model_list = []
model_dir = Path(shared.args.model_dir)

View file

@ -21,7 +21,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -34,6 +34,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -21,7 +21,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -34,6 +34,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -19,7 +19,7 @@ python-docx==1.1.2
pyyaml
requests
rich
safetensors==0.6.*
safetensors==0.7.*
scipy
sentencepiece
tensorboard
@ -32,6 +32,9 @@ wandb
gradio==4.37.*
https://github.com/oobabooga/gradio/releases/download/custom-build/gradio_client-1.0.2+custom.1-py3-none-any.whl
# Diffusers
diffusers @ git+https://github.com/huggingface/diffusers.git@edf36f5128abf3e6ecf92b5145115514363c58e6
# API
flask_cloudflared==0.0.14
sse-starlette==1.6.5

View file

@ -5,6 +5,7 @@ from pathlib import Path
from modules import shared
from modules.block_requests import OpenMonkeyPatch, RequestBlocker
from modules.image_models import load_image_model
from modules.logging_colors import logger
from modules.prompts import load_prompt
@ -50,6 +51,7 @@ from modules import (
ui_chat,
ui_default,
ui_file_saving,
ui_image_generation,
ui_model_menu,
ui_notebook,
ui_parameters,
@ -163,6 +165,7 @@ def create_interface():
ui_chat.create_character_settings_ui() # Character tab
ui_model_menu.create_ui() # Model tab
if not shared.args.portable:
ui_image_generation.create_ui() # Image generation tab
training.create_ui() # Training tab
ui_session.create_ui() # Session tab
@ -170,6 +173,8 @@ def create_interface():
ui_chat.create_event_handlers()
ui_default.create_event_handlers()
ui_notebook.create_event_handlers()
if not shared.args.portable:
ui_image_generation.create_event_handlers()
# Other events
ui_file_saving.create_event_handlers()
@ -256,6 +261,9 @@ if __name__ == "__main__":
if new_settings:
shared.settings.update(new_settings)
# Apply CLI overrides for image model settings (CLI flags take precedence over saved settings)
shared.apply_image_model_cli_overrides()
# Fallback settings for models
shared.model_config['.*'] = get_fallback_settings()
shared.model_config.move_to_end('.*', last=False) # Move to the beginning
@ -313,6 +321,22 @@ if __name__ == "__main__":
if shared.args.lora:
add_lora_to_model(shared.args.lora)
# Load image model if specified via CLI
if shared.args.image_model:
logger.info(f"Loading image model: {shared.args.image_model}")
result = load_image_model(
shared.args.image_model,
dtype=shared.settings.get('image_dtype', 'bfloat16'),
attn_backend=shared.settings.get('image_attn_backend', 'sdpa'),
cpu_offload=shared.settings.get('image_cpu_offload', False),
compile_model=shared.settings.get('image_compile', False),
quant_method=shared.settings.get('image_quant', 'none')
)
if result is not None:
shared.image_model_name = shared.args.image_model
else:
logger.error(f"Failed to load image model: {shared.args.image_model}")
shared.generation_lock = Lock()
if shared.args.idle_timeout > 0: