feifeilib/feifeitexttoimg.py

import gradio as gr
from PIL import Image
import numpy as np
import random
import spaces
import torch
import torch.nn.functional as F
import config
from extras.expansion import FooocusExpansion
import re

expansion = FooocusExpansion()
with open("artist.txt", "r") as file:
    artists = file.readlines()

MAX_SEED = np.iinfo(np.int32).max

# 去除每行末尾的换行符
artists = [artist.strip() for artist in artists]


@spaces.GPU()
def feifeitexttoimg(
    prompt,
    quality_select,
    sharpened_select,
    styles_Radio,
    FooocusExpansion_select,
    seed=42,
    randomize_seed=False,
    width=1024,
    height=1024,
    num_inference_steps=4,
    guidance_scale=3.5,
    num_strength=1,
    pipe=None,
    progress=gr.Progress(track_tqdm=True),
):
    guidance_scale = 3.5
    num_strength = 1
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)
    if not prompt:
        prompt = "the photo is a 18 yo jpop girl is looking absolutely adorable and gorgeous, with a playful and mischievous grin, her eyes twinkling with joy. art by __artist__ and __artist__"
    if "__artist__" in prompt:
        # 随机选择艺术家
        selected_artists = random.sample(artists, len(artists))

        # 使用正则表达式替换所有的 __artist__
        def replace_artists(match):
            return selected_artists.pop(0)

        prompt = re.sub(r"__artist__", replace_artists, prompt)
        # print("__artist__ " + prompt)
    if quality_select:
        prompt += ", masterpiece, best quality, very aesthetic, absurdres"
        # print("111 " + prompt)
    if styles_Radio:
        for style_name in styles_Radio:
            for style in config.style_list:
                if style["name"] == style_name:
                    prompt += style["prompt"].replace("{prompt}", "the ")
        # print("222 " + prompt)
    if FooocusExpansion_select:
        prompt = expansion(prompt, seed)
        # print("333 " + prompt)
    # print("000 " + prompt)

    # 生成图像
    image = pipe(
        prompt="",
        prompt_2=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale,
        output_type="pil",
    ).images[0]

    if sharpened_select:
        # 将PIL图像转换为NumPy数组
        image_np = np.array(image)

        # 将NumPy数组转换为PyTorch张量
        image_tensor = (
            torch.tensor(image_np).permute(2, 0, 1).unsqueeze(0).float().to("cuda")
        )

        # 定义锐化滤镜，并调整中心值
        strength = num_strength
        sharpen_kernel = (
            torch.tensor(
                [
                    [0, -1 * strength, 0],
                    [-1 * strength, 1 + 4 * strength, -1 * strength],
                    [0, -1 * strength, 0],
                ],
                dtype=torch.float32,
            )
            .unsqueeze(0)
            .unsqueeze(0)
            .to("cuda")
        )

        # 分别对每个通道应用卷积核
        sharpened_channels = []
        for i in range(3):
            channel_tensor = image_tensor[:, i : i + 1, :, :]
            sharpened_channel = F.conv2d(channel_tensor, sharpen_kernel, padding=1)
            sharpened_channels.append(sharpened_channel)

        # 合并通道
        sharpened_image_tensor = torch.cat(sharpened_channels, dim=1)

        # 将增强后的图像转换回PIL格式
        sharpened_image_np = (
            sharpened_image_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
        )
        sharpened_image_np = np.clip(sharpened_image_np, 0, 255).astype(np.uint8)
        image = Image.fromarray(sharpened_image_np)

    return image, seed