Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor api

Browse Source
This commit is contained in:
Vladimir Mandic
2024-01-28 10:12:24 -05:00
parent dc70246c78
commit 7e447222a1
13 changed files with 778 additions and 785 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/ui_interrogate.py
+36 -87
View File
@@ -1,11 +1,7 @@
import os
import base64
from io import BytesIO
import gradio as gr
import torch
from PIL import Image
from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
from fastapi.exceptions import HTTPException
import modules.generation_parameters_copypaste as parameters_copypaste
from modules import devices, lowvram, shared, paths, ui_common
@@ -29,7 +25,12 @@ class BatchWriter:
self.file.close()
def load(clip_model_name):
def get_models():
import open_clip
return ['/'.join(x) for x in open_clip.list_pretrained()]
def load_interrogator(clip_model_name):
from clip_interrogator import Config, Interrogator
global ci # pylint: disable=global-statement
if ci is None:
@@ -54,23 +55,6 @@ def unload():
devices.torch_gc()
def image_analysis(image, clip_model_name):
load(clip_model_name)
image = image.convert('RGB')
image_features = ci.image_to_features(image)
top_mediums = ci.mediums.rank(image_features, 5)
top_artists = ci.artists.rank(image_features, 5)
top_movements = ci.movements.rank(image_features, 5)
top_trendings = ci.trendings.rank(image_features, 5)
top_flavors = ci.flavors.rank(image_features, 5)
medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks
def interrogate(image, mode, caption=None):
shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
if mode == 'best':
@@ -88,14 +72,14 @@ def interrogate(image, mode, caption=None):
return prompt
def image_to_prompt(image, mode, clip_model_name):
def interrogate_image(image, model, mode):
shared.state.begin()
shared.state.job = 'interrogate'
try:
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu()
devices.torch_gc()
load(clip_model_name)
load_interrogator(model)
image = image.convert('RGB')
shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
prompt = interrogate(image, mode)
@@ -106,12 +90,7 @@ def image_to_prompt(image, mode, clip_model_name):
return prompt
def get_models():
import open_clip
return ['/'.join(x) for x in open_clip.list_pretrained()]
def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write):
def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
files = []
if batch_files is not None:
files += [f.name for f in batch_files]
@@ -122,7 +101,6 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
if len(files) == 0:
shared.log.error('Interrogate batch no images')
return ''
shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
shared.state.begin()
shared.state.job = 'batch interrogate'
prompts = []
@@ -130,7 +108,8 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu()
devices.torch_gc()
load(clip_model)
load_interrogator(model)
shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
captions = []
# first pass: generate captions
for file in files:
@@ -168,6 +147,23 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
return '\n\n'.join(prompts)
def analyze_image(image, model):
load_interrogator(model)
image = image.convert('RGB')
image_features = ci.image_to_features(image)
top_mediums = ci.mediums.rank(image_features, 5)
top_artists = ci.artists.rank(image_features, 5)
top_movements = ci.movements.rank(image_features, 5)
top_trendings = ci.trendings.rank(image_features, 5)
top_flavors = ci.flavors.rank(image_features, 5)
medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks
def create_ui():
global low_vram # pylint: disable=global-statement
low_vram = shared.cmd_opts.lowvram or shared.cmd_opts.medvram
@@ -190,9 +186,9 @@ def create_ui():
trending = gr.Label(label="Trending", num_top_classes=5)
flavor = gr.Label(label="Flavor", num_top_classes=5)
with gr.Row():
interrogate_btn = gr.Button("Interrogate", variant='primary')
analyze_btn = gr.Button("Analyze", variant='primary')
unload_btn = gr.Button("Unload")
btn_interrogate_img = gr.Button("Interrogate", variant='primary')
btn_analyze_img = gr.Button("Analyze", variant='primary')
btn_unload = gr.Button("Unload")
with gr.Row():
buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "extras", "control"])
for tabname, button in buttons.items():
@@ -209,7 +205,7 @@ def create_ui():
with gr.Row():
write = gr.Checkbox(label='Write prompts to files', value=False)
with gr.Row():
batch_btn = gr.Button("Interrogate", variant='primary')
btn_interrogate_batch = gr.Button("Interrogate", variant='primary')
with gr.Column():
with gr.Row():
# clip_model = gr.Dropdown(get_models(), value='ViT-L-14/openai', label='CLIP Model')
@@ -217,54 +213,7 @@ def create_ui():
ui_common.create_refresh_button(clip_model, get_models, lambda: {"choices": get_models()}, 'refresh_interrogate_models')
with gr.Row():
mode = gr.Radio(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='best')
interrogate_btn.click(image_to_prompt, inputs=[image, mode, clip_model], outputs=prompt)
analyze_btn.click(image_analysis, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
unload_btn.click(unload)
batch_btn.click(batch_process, inputs=[batch_files, batch_folder, batch_str, mode, clip_model, write], outputs=[batch])
def decode_base64_to_image(encoding):
if encoding.startswith("data:image/"):
encoding = encoding.split(";")[1].split(",")[1]
try:
image = Image.open(BytesIO(base64.b64decode(encoding)))
return image
except Exception as e:
raise HTTPException(status_code=500, detail="Invalid encoded image") from e
# TODO redesign interrogator api
def mount_interrogator_api(_: gr.Blocks, app):
class InterrogatorAnalyzeRequest(BaseModel):
image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
clip_model_name: str = Field(default="ViT-L-14/openai", title="Model", description="The interrogate model used. See the models endpoint for a list of available models.")
class InterrogatorPromptRequest(InterrogatorAnalyzeRequest):
mode: str = Field(default="fast", title="Mode", description="The mode used to generate the prompt. Can be one of: best, fast, classic, negative.")
@app.get("/interrogator/models")
async def api_get_models():
import open_clip
return ["/".join(x) for x in open_clip.list_pretrained()]
@app.post("/interrogator/prompt")
async def api_get_prompt(analyzereq: InterrogatorPromptRequest):
image_b64 = analyzereq.image
if image_b64 is None:
raise HTTPException(status_code=404, detail="Image not found")
img = decode_base64_to_image(image_b64)
prompt = image_to_prompt(img, analyzereq.mode, analyzereq.clip_model_name)
return {"prompt": prompt}
@app.post("/interrogator/analyze")
async def api_analyze(analyzereq: InterrogatorAnalyzeRequest):
image_b64 = analyzereq.image
if image_b64 is None:
raise HTTPException(status_code=404, detail="Image not found")
img = decode_base64_to_image(image_b64)
(medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks) = image_analysis(img, analyzereq.clip_model_name)
return {"medium": medium_ranks, "artist": artist_ranks, "movement": movement_ranks, "trending": trending_ranks, "flavor": flavor_ranks}
# script_callbacks.on_app_started(mount_interrogator_api)
btn_interrogate_img.click(interrogate_image, inputs=[image, clip_model, mode], outputs=prompt)
btn_analyze_img.click(analyze_image, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
btn_interrogate_batch.click(interrogate_batch, inputs=[batch_files, batch_folder, batch_str, clip_model, mode, write], outputs=[batch])
btn_unload.click(unload)