make generate mask from RemoveBG && AnimeSeg work

This commit is contained in:
Qing
2024-01-02 22:32:40 +08:00
parent 6253016019
commit aca85543ca
22 changed files with 244 additions and 100 deletions

View File

@@ -3,7 +3,7 @@
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>Lama Cleaner</title>
<title>IOPaint</title>
</head>
<body>
<div id="root"></div>

View File

@@ -272,7 +272,7 @@ export default function Editor(props: EditorProps) {
console.log("[useEffect] centerView")
// render 改变尺寸以后undo/redo 重新 center
viewportRef?.current?.centerView(minScale, 1)
}, [context?.canvas.height, context?.canvas.width, viewportRef, minScale])
}, [imageHeight, imageWidth, viewportRef, minScale])
// Zoom reset
const resetZoom = useCallback(() => {
@@ -358,6 +358,7 @@ export default function Editor(props: EditorProps) {
const targetFile = await getCurrentRender()
try {
const res = await runPlugin(
true,
PluginName.InteractiveSeg,
targetFile,
undefined,

View File

@@ -16,6 +16,7 @@ import {
Smile,
} from "lucide-react"
import { useStore } from "@/lib/states"
import { PluginInfo } from "@/lib/types"
export enum PluginName {
RemoveBG = "RemoveBG",
@@ -26,6 +27,7 @@ export enum PluginName {
InteractiveSeg = "InteractiveSeg",
}
// TODO: get plugin config from server and using form-render??
const pluginMap = {
[PluginName.RemoveBG]: {
IconClass: Slice,
@@ -37,7 +39,7 @@ const pluginMap = {
},
[PluginName.RealESRGAN]: {
IconClass: Fullscreen,
showName: "RealESRGAN 4x",
showName: "RealESRGAN",
},
[PluginName.GFPGAN]: {
IconClass: Smile,
@@ -67,11 +69,11 @@ const Plugins = () => {
return null
}
const onPluginClick = (pluginName: string) => {
const onPluginClick = (genMask: boolean, pluginName: string) => {
if (pluginName === PluginName.InteractiveSeg) {
updateInteractiveSegState({ isInteractiveSeg: true })
} else {
runRenderablePlugin(pluginName)
runRenderablePlugin(genMask, pluginName)
}
}
@@ -87,14 +89,14 @@ const Plugins = () => {
<DropdownMenuSubContent>
<DropdownMenuItem
onClick={() =>
runRenderablePlugin(PluginName.RealESRGAN, { upscale: 2 })
runRenderablePlugin(false, PluginName.RealESRGAN, { upscale: 2 })
}
>
upscale 2x
</DropdownMenuItem>
<DropdownMenuItem
onClick={() =>
runRenderablePlugin(PluginName.RealESRGAN, { upscale: 4 })
runRenderablePlugin(false, PluginName.RealESRGAN, { upscale: 4 })
}
>
upscale 4x
@@ -104,16 +106,44 @@ const Plugins = () => {
)
}
const renderGenImageAndMaskPlugin = (plugin: PluginInfo) => {
const { IconClass, showName } = pluginMap[plugin.name as PluginName]
return (
<DropdownMenuSub key={plugin.name}>
<DropdownMenuSubTrigger disabled={disabled}>
<div className="flex gap-2 items-center">
<IconClass className="p-1" />
{showName}
</div>
</DropdownMenuSubTrigger>
<DropdownMenuSubContent>
<DropdownMenuItem onClick={() => onPluginClick(false, plugin.name)}>
Remove Background
</DropdownMenuItem>
<DropdownMenuItem onClick={() => onPluginClick(true, plugin.name)}>
Generate Mask
</DropdownMenuItem>
</DropdownMenuSubContent>
</DropdownMenuSub>
)
}
const renderPlugins = () => {
return plugins.map((plugin: string) => {
const { IconClass, showName } = pluginMap[plugin as PluginName]
if (plugin === PluginName.RealESRGAN) {
return plugins.map((plugin: PluginInfo) => {
const { IconClass, showName } = pluginMap[plugin.name as PluginName]
if (plugin.name === PluginName.RealESRGAN) {
return renderRealESRGANPlugin()
}
if (
plugin.name === PluginName.RemoveBG ||
plugin.name === PluginName.AnimeSeg
) {
return renderGenImageAndMaskPlugin(plugin)
}
return (
<DropdownMenuItem
key={plugin}
onClick={() => onPluginClick(plugin)}
key={plugin.name}
onClick={() => onPluginClick(false, plugin.name)}
disabled={disabled}
>
<div className="flex gap-2 items-center">

View File

@@ -64,11 +64,6 @@ export function Shortcuts() {
<ShortCut content="Decrease Brush Size" keys={["["]} />
<ShortCut content="Increase Brush Size" keys={["]"]} />
<ShortCut content="View Original Image" keys={["Hold Tab"]} />
<ShortCut
content="Multi-Stroke Drawing"
keys={[`Hold ${CmdOrCtrl()}`]}
/>
<ShortCut content="Cancel Drawing" keys={["Esc"]} />
<ShortCut content="Undo" keys={[CmdOrCtrl(), "Z"]} />
<ShortCut content="Redo" keys={[CmdOrCtrl(), "Shift", "Z"]} />

View File

@@ -658,13 +658,13 @@ const DiffusionOptions = () => {
updateSettings({ sdSampler: value })
}}
>
<SelectTrigger className="w-[180px]">
<SelectTrigger className="w-[175px] text-xs">
<SelectValue placeholder="Select sampler" />
</SelectTrigger>
<SelectContent align="end">
<SelectGroup>
{samplers.map((sampler) => (
<SelectItem key={sampler} value={sampler}>
<SelectItem key={sampler} value={sampler} className="text-xs">
{sampler}
</SelectItem>
))}

View File

@@ -17,6 +17,7 @@ const ToastViewport = React.forwardRef<
"fixed top-0 z-[100] flex max-h-screen w-full flex-col-reverse p-4 sm:bottom-0 sm:right-0 sm:top-auto sm:flex-col md:max-w-[420px]",
className
)}
tabIndex={-1}
{...props}
/>
))
@@ -47,6 +48,7 @@ const Toast = React.forwardRef<
<ToastPrimitives.Root
ref={ref}
className={cn(toastVariants({ variant }), className)}
tabIndex={-1}
{...props}
/>
)
@@ -63,6 +65,7 @@ const ToastAction = React.forwardRef<
"inline-flex h-8 shrink-0 items-center justify-center rounded-md border bg-transparent px-3 text-sm font-medium transition-colors hover:bg-secondary focus:outline-none focus:ring-1 focus:ring-ring disabled:pointer-events-none disabled:opacity-50 group-[.destructive]:border-muted/40 group-[.destructive]:hover:border-destructive/30 group-[.destructive]:hover:bg-destructive group-[.destructive]:hover:text-destructive-foreground group-[.destructive]:focus:ring-destructive",
className
)}
tabIndex={-1}
{...props}
/>
))
@@ -79,6 +82,7 @@ const ToastClose = React.forwardRef<
className
)}
toast-close=""
tabIndex={-1}
{...props}
>
<Cross2Icon className="h-4 w-4" />
@@ -93,6 +97,7 @@ const ToastTitle = React.forwardRef<
<ToastPrimitives.Title
ref={ref}
className={cn("text-sm font-semibold [&+div]:text-xs", className)}
tabIndex={-1}
{...props}
/>
))
@@ -106,6 +111,7 @@ const ToastDescription = React.forwardRef<
ref={ref}
className={cn("text-sm opacity-90", className)}
{...props}
tabIndex={-1}
/>
))
ToastDescription.displayName = ToastPrimitives.Description.displayName

View File

@@ -114,13 +114,15 @@ export function fetchModelInfos(): Promise<ModelInfo[]> {
}
export async function runPlugin(
genMask: boolean,
name: string,
imageFile: File,
upscale?: number,
clicks?: number[][]
) {
const imageBase64 = await convertToBase64(imageFile)
const res = await fetch(`${API_ENDPOINT}/run_plugin`, {
const p = genMask ? "run_plugin_gen_mask" : "run_plugin_gen_image"
const res = await fetch(`${API_ENDPOINT}/${p}`, {
method: "POST",
headers: {
"Content-Type": "application/json",

View File

@@ -156,7 +156,6 @@ type AppAction = {
setFile: (file: File) => Promise<void>
setCustomFile: (file: File) => void
setIsInpainting: (newValue: boolean) => void
setIsPluginRunning: (newValue: boolean) => void
getIsProcessing: () => boolean
setBaseBrushSize: (newValue: number) => void
getBrushSize: () => number
@@ -190,6 +189,7 @@ type AppAction = {
showPrevMask: () => Promise<void>
hidePrevMask: () => void
runRenderablePlugin: (
genMask: boolean,
pluginName: string,
params?: PluginParams
) => Promise<void>
@@ -521,28 +521,43 @@ export const useStore = createWithEqualityFn<AppState & AppAction>()(
},
runRenderablePlugin: async (
genMask: boolean,
pluginName: string,
params: PluginParams = { upscale: 1 }
) => {
const { renders, lineGroups } = get().editorState
set((state) => {
state.isInpainting = true
state.isPluginRunning = true
})
try {
const start = new Date()
const targetFile = await get().getCurrentTargetFile()
const res = await runPlugin(pluginName, targetFile, params.upscale)
const res = await runPlugin(
genMask,
pluginName,
targetFile,
params.upscale
)
const { blob } = res
const newRender = new Image()
await loadImage(newRender, blob)
get().setImageSize(newRender.width, newRender.height)
const newRenders = [...renders, newRender]
const newLineGroups = [...lineGroups, []]
get().updateEditorState({
renders: newRenders,
lineGroups: newLineGroups,
})
if (!genMask) {
const newRender = new Image()
await loadImage(newRender, blob)
get().setImageSize(newRender.width, newRender.height)
const newRenders = [...renders, newRender]
const newLineGroups = [...lineGroups, []]
get().updateEditorState({
renders: newRenders,
lineGroups: newLineGroups,
})
} else {
const newMask = new Image()
await loadImage(newMask, blob)
get().updateInteractiveSegState({
interactiveSegMask: newMask,
})
}
const end = new Date()
const time = end.getTime() - start.getTime()
toast({
@@ -555,7 +570,7 @@ export const useStore = createWithEqualityFn<AppState & AppAction>()(
})
}
set((state) => {
state.isInpainting = false
state.isPluginRunning = false
})
},
@@ -803,11 +818,6 @@ export const useStore = createWithEqualityFn<AppState & AppAction>()(
state.isInpainting = newValue
}),
setIsPluginRunning: (newValue: boolean) =>
set((state) => {
state.isPluginRunning = newValue
}),
setFile: async (file: File) => {
if (get().settings.enableAutoExtractPrompt) {
try {

View File

@@ -6,8 +6,14 @@ export interface Filename {
mtime: number
}
export interface PluginInfo {
name: string
support_gen_image: boolean
support_gen_mask: boolean
}
export interface ServerConfig {
plugins: string[]
plugins: PluginInfo[]
enableFileManager: boolean
enableAutoSaving: boolean
enableControlnet: boolean