generative_front/components/tabs/textImagePanel.js

// components/PhoneInput.tsx
import {
  Button,
  Divider,
  FormControl,
  FormLabel,
  HStack,
  Icon,
  Input,
  Table,
  TableContainer,
  Tbody,
  Textarea,
  Th,
  Thead,
  Tr,
  VStack,
} from "@chakra-ui/react";
import { MdHeight } from "react-icons/md";
import { RxWidth } from "react-icons/rx";
import { BsBarChartSteps } from "react-icons/bs";
import { TbPrompt } from "react-icons/tb";
import { FaDraft2Digital, FaPlay } from "react-icons/fa";
import { useEffect, useState } from "react";
import { useForm } from "react-hook-form";
import ResultBox from "../ui/ResultBox";
import axiosInstance, { baseUrl } from "@/lib/api";
import useSWRMutation from "swr/mutation";

const modelsDic = {
  sdxl: "base",
  sdxl_lora: "lora",
};

const postRequest = async (url, { arg }) => {
  const response = await axiosInstance.post(baseUrl + url, arg);
  return response?.data;
};

export default function TextImagePanel({ filters }) {
  const { register, handleSubmit, watch, setValue } = useForm("");
  const {
    register: registerRokn,
    handleSubmit: handleSubmitRokn,
    watch: watchRokn,
    setValue: setValueRokn,
  } = useForm("");

  const [images, setImages] = useState([]);

  const width = filters?.model_name === 0 ? watch("width") : watchRokn("width");
  const height =
    filters?.model_name === 0 ? watch("height") : watchRokn("height");

  const { trigger: triggerGenerateImage, isMutating: isGeneratingImage } =
    useSWRMutation(
      `/content/sdxl/${modelsDic[filters?.tab_menu]}`,
      postRequest,
      {
        onSuccess: (data) => setImages(data?.images),
      },
    );

  const {
    trigger: triggerGenerateRoknImage,
    isMutating: isGeneratingRoknImage,
  } = useSWRMutation(
    `/content/rokn/t2i?model_key=${filters?.tab_menu}`,
    postRequest,
    {
      onSuccess: (data) => setImages(data?.images),
    },
  );

  const onSubmit = (data) => {
    if (filters?.model_name === 0) triggerGenerateImage(data);
    else triggerGenerateRoknImage(data);
  };

  useEffect(() => {
    if (filters?.model_name === 0) {
      setValue("height", 768);
      setValue("width", 768);
      setValue("seed", 42);
      setValue("num_inference_steps", 60);
      setValue("guidance_scale", 7);
      setValue("prompt", "");
    } else {
      setValueRokn("height", 768);
      setValueRokn("width", 768);
      setValueRokn("steps", 10);
      setValueRokn("preset", "");
      setValue("prompt", "");
    }
  }, [filters?.model_name]);

  return (
    <VStack
      spacing={4}
      w="100%"
      h="100%"
      alignItems="start"
      overflowY="auto"
      as={"form"}
      onSubmit={
        filters?.model_name === 0
          ? handleSubmit(onSubmit)
          : handleSubmitRokn(onSubmit)
      }
    >
      {" "}
      <HStack>
        <FormControl as={VStack} alignItems={"start"}>
          <HStack w={"100%"} alignItems={"start"}>
            <Icon
              as={MdHeight}
              bgColor={"gray.200"}
              borderRadius={"50%"}
              p={"5px"}
              fontSize={"25px"}
              border={"1px"}
              borderColor={"gray.300"}
            />
            <FormLabel>طول:</FormLabel>
          </HStack>
          {filters?.model_name === 0 ? (
            <Input bgColor={"white"} {...register("height")} type="number" />
          ) : (
            <Input
              bgColor={"white"}
              {...registerRokn("height")}
              type="number"
            />
          )}
        </FormControl>
        <FormControl as={VStack} alignItems={"start"}>
          <HStack w={"100%"} alignItems={"start"}>
            <Icon
              as={RxWidth}
              bgColor={"gray.200"}
              borderRadius={"50%"}
              p={"5px"}
              fontSize={"25px"}
              border={"1px"}
              borderColor={"gray.300"}
            />
            <FormLabel>عرض:</FormLabel>
          </HStack>
          {filters?.model_name === 0 ? (
            <Input bgColor={"white"} {...register("width")} type="number" />
          ) : (
            <Input bgColor={"white"} {...registerRokn("width")} type="number" />
          )}
        </FormControl>
        {filters?.model_name === 0 && (
          <FormControl as={VStack} alignItems={"start"}>
            <FormLabel>seed:</FormLabel>
            <Input bgColor={"white"} {...register("seed")} type="number" />
          </FormControl>
        )}
        <FormControl as={VStack} alignItems={"start"}>
          <HStack alignItems={"start"}>
            <Icon
              as={BsBarChartSteps}
              bgColor={"gray.200"}
              borderRadius={"50%"}
              p={"5px"}
              fontSize={"25px"}
              border={"1px"}
              borderColor={"gray.300"}
            />
            <FormLabel>تعداد مراحل استنتاج:</FormLabel>
          </HStack>
          {filters?.model_name === 0 ? (
            <Input
              bgColor={"white"}
              {...register("num_inference_steps")}
              type="number"
            />
          ) : (
            <Input bgColor={"white"} {...registerRokn("steps")} type="number" />
          )}
        </FormControl>
        {filters?.model_name === 0 && (
          <FormControl as={VStack} alignItems={"start"}>
            <HStack alignItems={"start"}>
              <Icon
                as={FaDraft2Digital}
                bgColor={"gray.200"}
                borderRadius={"50%"}
                p={"5px"}
                fontSize={"25px"}
                border={"1px"}
                borderColor={"gray.300"}
              />
              <FormLabel>مقیاس راهنمایی:</FormLabel>
            </HStack>
            <Input
              bgColor={"white"}
              {...register("guidance_scale")}
              type="number"
            />
          </FormControl>
        )}
        {filters?.model_name === 1 && (
          <FormControl as={VStack} alignItems={"start"}>
            <HStack alignItems={"start"}>
              <Icon
                as={FaDraft2Digital}
                bgColor={"gray.200"}
                borderRadius={"50%"}
                p={"5px"}
                fontSize={"25px"}
                border={"1px"}
                borderColor={"gray.300"}
              />
              <FormLabel>سبک تصویر:</FormLabel>
            </HStack>
            <Input bgColor={"white"} {...registerRokn("preset")} />
          </FormControl>
        )}
      </HStack>
      <FormControl as={VStack} alignItems={"start"} w={"auto"}>
        <HStack>
          <Icon
            as={FaDraft2Digital}
            bgColor={"gray.200"}
            borderRadius={"50%"}
            p={"5px"}
            fontSize={"25px"}
            border={"1px"}
            borderColor={"gray.300"}
          />
          <FormLabel>تعداد:</FormLabel>
        </HStack>
        {filters?.model_name === 0 ? (
          <Input bgColor={"white"} {...register("number")} type="number" />
        ) : (
          <Input bgColor={"white"} {...registerRokn("number")} type="number" />
        )}
      </FormControl>
      <FormControl as={VStack} alignItems={"start"}>
        <HStack alignItems={"start"}>
          <Icon
            as={TbPrompt}
            bgColor={"gray.200"}
            borderRadius={"50%"}
            p={"5px"}
            fontSize={"25px"}
            border={"1px"}
            borderColor={"gray.300"}
          />
          <FormLabel>پرامپت:</FormLabel>
        </HStack>
        {filters?.model_name === 0 ? (
          <Textarea bgColor={"white"} {...register("prompt")} />
        ) : (
          <Textarea bgColor={"white"} {...registerRokn("prompt")} />
        )}
      </FormControl>
      <Button
        p={"20px"}
        colorScheme="pink"
        leftIcon={<FaPlay />}
        type="submit"
        isLoading={isGeneratingImage || isGeneratingRoknImage}
        alignSelf={"end"}
      >
        اجرا
      </Button>
      {/* <Divider my={"20px"} borderColor={"gray.300"} />
      <TableContainer w={"100%"} mt={"10px"}>
        <Table>
          <Thead>
            <Tr>
              <Th fontSize={"md"}>ردیف</Th>
            </Tr>
          </Thead>
          <Tbody></Tbody>
        </Table>
      </TableContainer> */}
      <Divider my={"20px"} borderColor={"gray.300"} />
      <ResultBox
        type="image"
        loading={isGeneratingImage || isGeneratingRoknImage}
        resultData={images}
        imgWidth={width || 500}
        imgHeight={height || 500}
      />
    </VStack>
  );
}