test_glm

app.py

@@ -1,4 +1,5 @@
 import os
+import re
 import threading
 import time
 from datetime import datetime, timedelta
@@ -9,6 +10,7 @@ import random
 import spaces  # [uncomment to use ZeroGPU]
 from diffusers import FluxPipeline
 import torch
+from openai import OpenAI
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "black-forest-labs/FLUX.1-dev"
@@ -24,10 +26,95 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 
 
+def clean_string(s):
+    s = s.replace("\n", " ")
+    s = s.strip()
+    s = re.sub(r"\s{2,}", " ", s)
+    return s
+
+
+def convert_prompt(
+        prompt: str,
+        retry_times: int = 5,
+) -> str:
+    if not os.environ.get("OPENAI_API_KEY"):
+        return prompt
+    client = OpenAI()
+    system_instruction = """
+    You are part of a team of bots that creates images . You work with an assistant bot that will draw anything you say. 
+    For example , outputting " a beautiful morning in the woods with the sun peaking through the trees " will trigger your partner bot to output an image of a forest morning , as described. 
+    You will be prompted by people looking to create detailed , amazing images. The way to accomplish this is to take their short prompts and make them extremely detailed and descriptive. 
+    There are a few rules to follow : 
+    - Prompt should always be written in English, regardless of the input language. Please provide the prompts in English.
+    - You will only ever output a single image description per user request.
+    - Image descriptions must be detailed and specific, including keyword categories such as subject, medium, style, additional details, color, and lighting. 
+    - When generating descriptions, focus on portraying the visual elements rather than delving into abstract psychological and emotional aspects. Provide clear and concise details that vividly depict the scene and its composition, capturing the tangible elements that make up the setting.
+    - Do not provide the process and explanation, just return the modified English description . Image descriptions must be between 100-200 words. Extra words will be ignored. 
+    """
+
+    text = prompt.strip()
+    for i in range(retry_times):
+        try:
+            response = client.chat.completions.create(
+                messages=[
+                    {"role": "system", "content": f"{system_instruction}"},
+                    {
+                        "role": "user",
+                        "content": 'Create an imaginative image descriptive caption for the user input : "一个头发花白的老人"',
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "A seasoned male with white hair and a neatly groomed beard stands confidently, donning a dark vest over a striped shirt. His hands are clasped together in front, one adorned with a ring, as he looks directly at the viewer with a composed expression. The soft lighting accentuates his features and the subtle textures of his attire, creating a portrait that exudes sophistication and a timeless elegance.",
+                    },
+                    {
+                        "role": "user",
+                        "content": 'Create an imaginative image descriptive caption for the user input : "画一只老鹰"',
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "A majestic eagle with expansive brown and white wings glides through the air, its sharp yellow eyes focused intently ahead. The eagle's talons are poised and ready for hunting, as it soars over a rugged mountainous terrain dusted with snow, under a soft blue sky.",
+                    },
+                    {
+                        "role": "user",
+                        "content": 'Create an imaginative image descriptive caption for the user input : "画一辆摩托车"',
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "Parked on a wet city street at night, a sleek motorcycle with a black and green design stands out. Its headlights cast a soft glow, reflecting off the puddles and highlighting its aerodynamic shape. The design is marked by sharp lines and angular features, with gold accents that shine against the dark backdrop. The motorcycle exudes an air of performance and luxury, ready to slice through the urban landscape.",
+                    },
+                    {
+                        "role": "user",
+                        "content": 'Create an imaginative image descriptive caption for the user input : "穿着金色盔甲的人"',
+                    },
+                    {
+                        "role": "assistant",
+                        "content": "A figure clad in meticulously crafted, golden armor stands with an air of quiet confidence. The armor, reminiscent of medieval knight attire, features a scalloped design with leaf-like patterns and is complemented by a black, form-fitting undergarment. The helmet, with its angular visor, adds to the intimidating presence. This armor, with its rich gold tones and intricate details, suggests a character of nobility or mythical origin, poised for valorous endeavors.",
+                    },
+                    {
+                        "role": "user",
+                        "content": f'Create an imaginative image descriptive caption for the user input : "{text}"',
+                    },
+                ],
+                model="glm-4-plus",
+                temperature=0.01,
+                top_p=0.7,
+                stream=False,
+                max_tokens=300,
+            )
+            prompt = response.choices[0].message.content
+            if prompt:
+                prompt = clean_string(prompt)
+                break
+        except Exception as e:
+            pass
+
+    return prompt
+
+
 def delete_old_files():
     while True:
         now = datetime.now()
-        cutoff = now - timedelta(minutes=10)
+        cutoff = now - timedelta(minutes=5)
         directories = ["./gradio_tmp"]
 
         for directory in directories:
@@ -102,12 +189,18 @@ with gr.Blocks(css=css) as demo:
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
-                max_lines=1,
+                max_lines=3,
                 placeholder="Enter your prompt",
                 container=False,
             )
-
-            run_button = gr.Button("Run", scale=0)
+        with gr.Row():
+            enhance = gr.Button("Enhance Prompt (Strongly Suggest)", scale=1)
+            enhance.click(
+                convert_prompt,
+                inputs=[prompt],
+                outputs=[prompt]
+            )
+            run_button = gr.Button("Run", scale=1)
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):

requirements.txt

@@ -3,4 +3,5 @@ diffusers>=0.30.3
 torch
 gradio>=4.44
 transformers>=4.45.2
-xformers
+xformers
+openai