Update feifeilib/feifeitexttoimg.py

feifeilib/feifeitexttoimg.py

@@ -1,121 +1,124 @@
-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
-from feifeilib.feifeimodload import feifeimodload
-import re
-
-expansion = FooocusExpansion()
-pipe = feifeimodload()
-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=False,
-        sharpened_select=False,
-        styles_Radio=["(None)"],
-        FooocusExpansion_select=False,
-        seed=random.randint(0, MAX_SEED),
-        randomize_seed=False,
-        width=896,
-        height=1152,
-        num_inference_steps=4,
-        guidance_scale=3.5,
-        num_strength=0.35,
-        progress=gr.Progress(track_tqdm=True),
-):
-    # 处理随机种子
-    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
+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
+from feifeilib.feifeimodload import feifeimodload
+import re
+from diffusers.callbacks import SDXLCFGCutoffCallback
+
+callback = SDXLCFGCutoffCallback(cutoff_step_ratio=0.4)
+expansion = FooocusExpansion()
+pipe = feifeimodload()
+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=False,
+        sharpened_select=False,
+        styles_Radio=["(None)"],
+        FooocusExpansion_select=False,
+        seed=random.randint(0, MAX_SEED),
+        randomize_seed=False,
+        width=896,
+        height=1152,
+        num_inference_steps=4,
+        guidance_scale=3.5,
+        num_strength=0.35,
+        progress=gr.Progress(track_tqdm=True),
+):
+    # 处理随机种子
+    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,
+        callback_on_step_end=callback,
+        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