From 90efc75d02056ff1efb073ee70afbe52af3426a7 Mon Sep 17 00:00:00 2001
From: BFlat <bombayapirat@gmail.com>
Date: Tue, 13 Feb 2024 15:39:29 +0700
Subject: [PATCH] image gen lora weight page update

---
 .../api/image_generation/EnqueueImageGen.ts   |  2 +-
 .../common/SelectModal/SelectModal.tsx        |  2 -
 .../src/v2/view/pages/weight/WeightPage.tsx   | 36 ++++++--
 .../inference_panels/SdInferencePanel.tsx     | 85 +++++++++++--------
 4 files changed, 82 insertions(+), 43 deletions(-)

diff --git a/packages/components/src/api/image_generation/EnqueueImageGen.ts b/packages/components/src/api/image_generation/EnqueueImageGen.ts
index b3400df405..b3859b6741 100644
--- a/packages/components/src/api/image_generation/EnqueueImageGen.ts
+++ b/packages/components/src/api/image_generation/EnqueueImageGen.ts
@@ -2,7 +2,7 @@ import { ApiConfig } from "../ApiConfig";
 
 export interface EnqueueImageGenRequest {
   uuid_idempotency_token: string;
-  maybe_sd_model_token: string | undefined;
+  maybe_sd_model_token: string | null;
   maybe_lora_model_token: string | null;
   maybe_prompt: string;
   maybe_n_prompt: string;
diff --git a/packages/fakeyou.com/src/components/common/SelectModal/SelectModal.tsx b/packages/fakeyou.com/src/components/common/SelectModal/SelectModal.tsx
index c13a9d193c..0824caa530 100644
--- a/packages/fakeyou.com/src/components/common/SelectModal/SelectModal.tsx
+++ b/packages/fakeyou.com/src/components/common/SelectModal/SelectModal.tsx
@@ -46,8 +46,6 @@ const SelectModal = memo(
         activeTab: tabs[0].tabKey,
         valueType: tabs[0].typeFilter || "all",
       });
-    console.log("🔫", tabs, activeTab, valueType);
-
     // Update mediaType when activeTab changes
     useEffect(() => {
       const currentTab = tabs.find(tab => tab.tabKey === activeTab);
diff --git a/packages/fakeyou.com/src/v2/view/pages/weight/WeightPage.tsx b/packages/fakeyou.com/src/v2/view/pages/weight/WeightPage.tsx
index dd1d7f65f1..c4f67d179a 100644
--- a/packages/fakeyou.com/src/v2/view/pages/weight/WeightPage.tsx
+++ b/packages/fakeyou.com/src/v2/view/pages/weight/WeightPage.tsx
@@ -110,6 +110,7 @@ export default function WeightPage({
 
   const deleteWeight = () => remove(!!user?.can_ban_users);
 
+  //Studio Access feature flag
   switch (weight?.weight_type) {
     case WeightType.SD_15:
     case WeightType.SDXL:
@@ -119,6 +120,34 @@ export default function WeightPage({
       }
   }
 
+  //Image generation panel if it's a lora weight or sd weight
+  let imageGenPanel = <></>;
+  switch (weight?.weight_type) {
+    case WeightType.SD_15:
+    case WeightType.SDXL:
+      imageGenPanel = (
+        <SdInferencePanel
+          inferenceJobs={inferenceJobs}
+          sessionSubscriptionsWrapper={sessionSubscriptionsWrapper}
+          weight_token={weight?.weight_token}
+          enqueueInferenceJob={enqueueInferenceJob}
+          weightPageType="sd"
+        />
+      );
+      break;
+    case WeightType.LORA:
+      imageGenPanel = (
+        <SdInferencePanel
+          inferenceJobs={inferenceJobs}
+          sessionSubscriptionsWrapper={sessionSubscriptionsWrapper}
+          weight_token={weight?.weight_token}
+          enqueueInferenceJob={enqueueInferenceJob}
+          weightPageType="lora"
+        />
+      );
+      break;
+  }
+
   function renderWeightComponent(weight: Weight) {
     switch (weight.weight_category) {
       case WeightCategory.TTS:
@@ -168,12 +197,7 @@ export default function WeightPage({
         return (
           <div className="d-flex flex-column gap-3">
             <SdCoverImagePanel src={sdCoverImage} />
-            <SdInferencePanel
-              inferenceJobs={inferenceJobs}
-              sessionSubscriptionsWrapper={sessionSubscriptionsWrapper}
-              sd_model_token={weight.weight_token}
-              enqueueInferenceJob={enqueueInferenceJob}
-            />
+            {imageGenPanel}
           </div>
         );
       default:
diff --git a/packages/fakeyou.com/src/v2/view/pages/weight/inference_panels/SdInferencePanel.tsx b/packages/fakeyou.com/src/v2/view/pages/weight/inference_panels/SdInferencePanel.tsx
index e81515fde9..a9421e5cab 100644
--- a/packages/fakeyou.com/src/v2/view/pages/weight/inference_panels/SdInferencePanel.tsx
+++ b/packages/fakeyou.com/src/v2/view/pages/weight/inference_panels/SdInferencePanel.tsx
@@ -1,4 +1,4 @@
-import React, { useRef, useState, memo } from "react";
+import React, { useRef, useState, memo, useEffect } from "react";
 import {
   Accordion,
   Button,
@@ -38,26 +38,38 @@ import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 interface SdInferencePanelProps {
   inferenceJobs: Array<InferenceJob>;
   sessionSubscriptionsWrapper: SessionSubscriptionsWrapper;
-  sd_model_token?: string;
+  weight_token?: string;
   enqueueInferenceJob: (
     jobToken: string,
     frontendInferenceJobType: FrontendInferenceJobType
   ) => void;
   isStandalone?: boolean;
+  weightPageType?: "lora" | "sd";
 }
 
 function SdInferencePanel({
   enqueueInferenceJob,
-  sd_model_token,
+  weight_token,
   sessionSubscriptionsWrapper,
   isStandalone = false,
   inferenceJobs,
+  weightPageType = "sd",
 }: SdInferencePanelProps) {
   const [loraToken, setLoraToken] = useState<string | null>(null);
-  const [weightToken, setWeightToken] = useState(sd_model_token);
+  const [sdToken, setSdToken] = useState<string | null>(null);
+
+  useEffect(() => {
+    if (weightPageType === "sd") {
+      setSdToken(weight_token || null);
+      setLoraToken(null);
+    } else if (weightPageType === "lora") {
+      setLoraToken(weight_token || null);
+      setSdToken(null);
+    }
+  }, [weightPageType, weight_token]);
 
   const handleOnWeightSelect = (data: { token: string; title: string }) => {
-    setWeightToken(data.token);
+    setSdToken(data.token);
   };
   const handleOnSelect = (data: { token: string; title: string }) => {
     setLoraToken(data.token);
@@ -243,7 +255,7 @@ function SdInferencePanel({
   ) => {
     ev.preventDefault();
 
-    if (!prompt || weightToken === undefined) {
+    if (!prompt || sdToken === undefined) {
       return false;
     }
 
@@ -258,10 +270,12 @@ function SdInferencePanel({
       internalSeed.current = generateRandomSeed();
     }
 
+    console.log(sdToken);
+
     const request = {
       uuid_idempotency_token: uuidv4(),
-      maybe_sd_model_token: isStandalone ? weightToken : sd_model_token,
-      maybe_lora_model_token: loraToken,
+      maybe_sd_model_token: sdToken || null,
+      maybe_lora_model_token: loraToken || null,
       maybe_prompt: prompt,
       maybe_n_prompt: negativePrompt,
       maybe_seed: internalSeed.current,
@@ -323,11 +337,11 @@ function SdInferencePanel({
 
         <SplitPanel.Body padding={true}>
           <div className="d-flex flex-column gap-3 mb-4">
-            {isStandalone && (
+            {(isStandalone || weightPageType === "lora") && (
               <SelectModal
                 required={true}
-                modalTitle="Select a Weight"
-                label="Select a Weight"
+                modalTitle="Select Stable Diffusion Weight"
+                label="Select a Stable Diffusion Weight"
                 onSelect={handleOnWeightSelect}
                 tabs={[
                   {
@@ -433,27 +447,30 @@ function SdInferencePanel({
                     value: samples,
                   }}
                 />
-                <SelectModal
-                  label="Additional LoRA Weight"
-                  onSelect={handleOnSelect}
-                  tabs={[
-                    {
-                      label: "All LoRA Weights",
-                      tabKey: "allLoraWeights",
-                      typeFilter: "loRA",
-                      searcher: true,
-                      type: "weights",
-                    },
-                    {
-                      label: "Bookmarked",
-                      tabKey: "bookmarkedLoraWeights",
-                      typeFilter: "loRA",
-                      searcher: false,
-                      type: "weights",
-                      onlyBookmarked: true,
-                    },
-                  ]}
-                />
+                {weightPageType === "sd" && (
+                  <SelectModal
+                    modalTitle="Select LoRA Weight"
+                    label="Additional LoRA Weight"
+                    onSelect={handleOnSelect}
+                    tabs={[
+                      {
+                        label: "All LoRA Weights",
+                        tabKey: "allLoraWeights",
+                        typeFilter: "loRA",
+                        searcher: true,
+                        type: "weights",
+                      },
+                      {
+                        label: "Bookmarked",
+                        tabKey: "bookmarkedLoraWeights",
+                        typeFilter: "loRA",
+                        searcher: false,
+                        type: "weights",
+                        onlyBookmarked: true,
+                      },
+                    ]}
+                  />
+                )}
                 <div>
                   <Label label="Number of Generations" />
                   <PremiumLock
@@ -476,7 +493,7 @@ function SdInferencePanel({
           <div className="d-flex gap-2 justify-content-end mt-4 mb-3">
             <Button
               {...{
-                label: "Clear All",
+                label: "Clear/Reset ",
                 variant: "secondary",
                 onClick: resetToInitialState,
               }}
@@ -484,7 +501,7 @@ function SdInferencePanel({
             <Button
               {...{
                 label: "Generate Image",
-                disabled: prompt === "" || weightToken === "",
+                disabled: prompt === "" || sdToken === "",
                 onClick: handleEnqueueImageGen,
                 isLoading: isEnqueuing,
               }}