From 774b5975736ba3f750784f256239db0bb292ca58 Mon Sep 17 00:00:00 2001
From: kwaroran <amasio1234@proton.me>
Date: Fri, 31 May 2024 13:31:02 +0900
Subject: [PATCH] refactor: Add PlaygroundImageGen component to PlaygroundMenu

---
 src/lib/Playground/PlaygroundImageGen.svelte | 49 ++++++++++++++++++++
 src/lib/Playground/PlaygroundMenu.svelte     |  9 ++++
 src/ts/process/stableDiff.ts                 | 38 ++++++++++++---
 3 files changed, 90 insertions(+), 6 deletions(-)
 create mode 100644 src/lib/Playground/PlaygroundImageGen.svelte

diff --git a/src/lib/Playground/PlaygroundImageGen.svelte b/src/lib/Playground/PlaygroundImageGen.svelte
new file mode 100644
index 00000000..53f71a19
--- /dev/null
+++ b/src/lib/Playground/PlaygroundImageGen.svelte
@@ -0,0 +1,49 @@
+<script lang="ts">
+    import { language } from "src/lang";
+    import TextAreaInput from "../UI/GUI/TextAreaInput.svelte";
+    import Button from "../UI/GUI/Button.svelte";
+    import { generateAIImage } from "src/ts/process/stableDiff";
+    import { createBlankChar } from "src/ts/characters";
+
+    let prompt = "";
+    let negPrompt = "";
+    let img = "";
+    let generating = false
+
+    const run = async () => {
+        console.log('running')
+        if(generating){
+            return
+        }
+        generating = true
+        const gen = await generateAIImage(prompt, createBlankChar(), negPrompt, 'inlay')
+        generating = false
+        if(gen){
+            img = gen
+        }
+    }
+</script>
+
+<h2 class="text-4xl text-textcolor my-6 font-black relative">{language.imageGeneration}</h2>
+
+<span class="text-textcolor text-lg">Prompt</span>
+
+<TextAreaInput bind:value={prompt} />
+
+<span class="text-textcolor text-lg">Neg. Prompt</span>
+
+<TextAreaInput bind:value={negPrompt} />
+
+{#if img}
+    <span class="text-textcolor text-lg">Generated</span>
+
+    <img src={img} class="max-w-full mt-4" alt="Generated"/>
+
+{/if}
+<Button className="mt-6" on:click={run}>
+    {#if generating}
+        <div class="loadmove" />
+    {:else}
+        Generate
+    {/if}
+</Button>
\ No newline at end of file
diff --git a/src/lib/Playground/PlaygroundMenu.svelte b/src/lib/Playground/PlaygroundMenu.svelte
index c538af1d..cb3ac70e 100644
--- a/src/lib/Playground/PlaygroundMenu.svelte
+++ b/src/lib/Playground/PlaygroundMenu.svelte
@@ -10,6 +10,7 @@
     import { characterFormatUpdate, createBlankChar } from "src/ts/characters";
     import { get } from "svelte/store";
     import { DataBase, setDatabase, type character } from "src/ts/storage/database";
+  import PlaygroundImageGen from "./PlaygroundImageGen.svelte";
 
     const playgroundChat = () => {
         let db = get(DataBase)
@@ -69,6 +70,11 @@
             }}>
                 <h1 class="text-2xl font-bold text-start">Jinja</h1>
             </button>
+            <button class="bg-darkbg rounded-md p-6 flex flex-col transition-shadow hover:ring-1" on:click={() => {
+                PlaygroundStore.set(7)
+            }}>
+                <h1 class="text-2xl font-bold text-start">{language.imageGeneration}</h1>
+            </button>
         </div>
     {:else}
         {#if $SizeStore.w < 1024}
@@ -96,6 +102,9 @@
             {#if $PlaygroundStore === 6}
                 <PlaygroundJinja/>
             {/if}
+            {#if $PlaygroundStore === 7}
+                <PlaygroundImageGen/>
+            {/if}  
         </div>
     {/if}
 </div>
\ No newline at end of file
diff --git a/src/ts/process/stableDiff.ts b/src/ts/process/stableDiff.ts
index 18125aef..f52902a7 100644
--- a/src/ts/process/stableDiff.ts
+++ b/src/ts/process/stableDiff.ts
@@ -54,7 +54,7 @@ export async function stableDiff(currentChar:character,prompt:string){
     return await generateAIImage(genPrompt, currentChar, neg, '')
 }
 
-export async function generateAIImage(genPrompt:string, currentChar:character, neg:string, returnSdData:string){
+export async function generateAIImage(genPrompt:string, currentChar:character, neg:string, returnSdData:string):Promise<string|false>{
     const db = get(DataBase)
     if(db.sdProvider === 'webui'){
 
@@ -91,8 +91,7 @@ export async function generateAIImage(genPrompt:string, currentChar:character, n
                     return ''
                 }
             }
-
-            if(da.ok){
+            else if(da.ok){
                 let charemotions = get(CharEmotion)
                 const img = `data:image/png;base64,${da.data.images[0]}`
                 console.log(img)
@@ -300,7 +299,7 @@ export async function generateAIImage(genPrompt:string, currentChar:character, n
                 }
             }
 
-            if(da.ok){
+            else if(da.ok){
                 let charemotions = get(CharEmotion)
                 const img = await processZip(da.data);
                 const emos:[string, string,number][] = [[img, img, Date.now()]]
@@ -321,7 +320,7 @@ export async function generateAIImage(genPrompt:string, currentChar:character, n
         }
     }
     if(db.sdProvider === 'dalle'){
-        const data = await globalFetch("https://api.openai.com/v1/images/generations", {
+        const da = await globalFetch("https://api.openai.com/v1/images/generations", {
             body: {
                 "prompt": genPrompt,
                 "model": "dall-e-3",
@@ -333,7 +332,34 @@ export async function generateAIImage(genPrompt:string, currentChar:character, n
             }
         })
 
-        return data?.data?.url ?? ""
+        console.log(da)
+
+        if(returnSdData === 'inlay'){
+            let res = da?.data?.data?.[0]?.b64_json
+            if(!res){
+                alertError(JSON.stringify(da.data))
+                return ''
+            }
+            return `data:image/png;base64,${res}`
+        }
+
+        else if(da.ok){
+            let charemotions = get(CharEmotion)
+            let img = da?.data?.data?.[0]?.b64_json
+            if(!img){
+                alertError(JSON.stringify(da.data))
+                return false
+            }
+            img = `data:image/png;base64,${img}`
+            const emos:[string, string,number][] = [[img, img, Date.now()]]
+            charemotions[currentChar.chaId] = emos
+            CharEmotion.set(charemotions)
+        }
+        else{
+            alertError(Buffer.from(da.data).toString())
+            return false   
+        }
+        return returnSdData
     }
     return ''
 }
\ No newline at end of file