Update app.py

app.py

@@ -14,10 +14,11 @@ from kolors.models.controlnet import ControlNetModel
 from diffusers import AutoencoderKL
 from kolors.models.unet_2d_condition import UNet2DConditionModel
 from diffusers import EulerDiscreteScheduler
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from annotator.midas import MidasDetector
 from annotator.dwpose import DWposeDetector
 from annotator.util import resize_image, HWC3
+import os
 
 device = "cuda"
 ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
@@ -253,6 +254,59 @@ def load_description(fp):
         content = f.read()
     return content
 
+# Add the text_to_image function
+def text_to_image(text, size, position):
+    width, height = 1024, 576
+    image = Image.new("RGB", (width, height), "white")
+    draw = ImageDraw.Draw(image)
+    
+    font_files = ["Arial_Unicode.ttf"]
+    font = None
+    for font_file in font_files:
+        font_path = os.path.join(os.path.dirname(__file__), font_file)
+        if os.path.exists(font_path):
+            try:
+                font = ImageFont.truetype(font_path, size=size)
+                print(f"Using font: {font_file}")
+                break
+            except IOError:
+                print(f"Error loading font: {font_file}")
+    if font is None:
+        print("No suitable font found. Using default font.")
+        font = ImageFont.load_default()
+
+    lines = text.split('\n')
+    max_line_width = 0
+    total_height = 0
+    line_heights = []
+    for line in lines:
+        left, top, right, bottom = draw.textbbox((0, 0), line, font=font)
+        line_width = right - left
+        line_height = bottom - top
+        line_heights.append(line_height)
+        max_line_width = max(max_line_width, line_width)
+        total_height += line_height
+
+    position_mapping = {
+        "top-left": (10, 10),
+        "top-center": ((width - max_line_width) / 2, 10),
+        "top-right": (width - max_line_width - 10, 10),
+        "middle-left": (10, (height - total_height) / 2),
+        "middle-center": ((width - max_line_width) / 2, (height - total_height) / 2),
+        "middle-right": (width - max_line_width - 10, (height - total_height) / 2),
+        "bottom-left": (10, height - total_height - 10),
+        "bottom-center": ((width - max_line_width) / 2, height - total_height - 10),
+        "bottom-right": (width - max_line_width - 10, height - total_height - 10),
+    }
+
+    x, y = position_mapping.get(position, ((width - max_line_width) / 2, height - total_height - 10))
+    for i, line in enumerate(lines):
+        draw.text((x, y), line, fill="black", font=font)
+        y += line_heights[i]
+
+    return image
+
+# Modify the main Gradio interface
 with gr.Blocks(css=css) as Kolors:
     with gr.Row():
         with gr.Column(elem_id="col-left"):
@@ -263,7 +317,39 @@ with gr.Blocks(css=css) as Kolors:
                     lines=2
                 )
             with gr.Row():
-                image = gr.Image(label="Image", type="pil")
+                image_input_type = gr.Radio(["Upload Image", "Generate Text Image"], label="Input Type", value="Upload Image")
+            
+            with gr.Row():
+                image = gr.Image(label="Image", type="pil", visible=True)
+                with gr.Column(visible=False) as text_image_inputs:
+                    text_input = gr.Textbox(label="Enter Text", lines=5, placeholder="Type your text here...")
+                    font_size = gr.Radio([48, 72, 96, 144], label="Font Size", value=72)
+                    text_position = gr.Dropdown(
+                        ["top-left", "top-center", "top-right", "middle-left", "middle-center", "middle-right", "bottom-left", "bottom-center", "bottom-right"],
+                        label="Text Position",
+                        value="middle-center"
+                    )
+                    generate_text_image = gr.Button("Generate Text Image")
+                    
+
+    def toggle_image_input(choice):
+        return {
+            image: gr.update(visible=choice == "Upload Image"),
+            text_image_inputs: gr.update(visible=choice == "Generate Text Image")
+        }
+
+    image_input_type.change(toggle_image_input, image_input_type, [image, text_image_inputs])
+
+    def generate_and_use_text_image(text, size, position):
+        text_image = text_to_image(text, size, position)
+        return text_image
+
+    generate_text_image.click(
+        generate_and_use_text_image,
+        inputs=[text_input, font_size, text_position],
+        outputs=image
+    )
+    
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt = gr.Textbox(
                     label="Negative prompt",