chore: lpw, image size

app.py

@@ -1,11 +1,16 @@
+import spaces
+
 import os
-import gradio as gr
-import numpy as np
 import random
-import spaces
+import math
+
+import torch
+import numpy as np
+
 from diffusers import DiffusionPipeline
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+
+import gradio as gr
 
 try:
     from dotenv import load_dotenv
@@ -38,12 +43,16 @@ dart = dart.requires_grad_(False)
 dart = torch.compile(dart)
 tokenizer = AutoTokenizer.from_pretrained(DART_V3_REPO_ID)
 
-pipe = DiffusionPipeline.from_pretrained(IMAGE_MODEL_REPO_ID, torch_dtype=torch_dtype)
+pipe = DiffusionPipeline.from_pretrained(
+    IMAGE_MODEL_REPO_ID, 
+    torch_dtype=torch_dtype,
+    custom_pipeline="lpw_stable_diffusion_xl"
+)
 pipe = pipe.to(device)
 
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+MAX_IMAGE_SIZE = 2048
 
 TEMPLATE = (
     "<|bos|>"
@@ -59,6 +68,20 @@ TEMPLATE = (
     "<general>"
 )
 
+def get_aspect_ratio(width: int, height: int) -> str:
+    ar = width / height
+
+    if ar <= 1 / math.sqrt(3):
+        return "<|aspect_ratio:ultra_wide|>"
+    elif ar <= 8 / 9:  #
+        return "<|aspect_ratio:wide|>"
+    elif ar < 9 / 8:
+        return "<|aspect_ratio:square|>"
+    elif ar < math.sqrt(3):
+        return "<|aspect_ratio:tall|>"
+    else:
+        return "<|aspect_ratio:ultra_tall|>"
+
 
 @torch.inference_mode
 def generate_prompt(aspect_ratio: str):
@@ -66,7 +89,7 @@ def generate_prompt(aspect_ratio: str):
         TEMPLATE.format(aspect_ratio=aspect_ratio),
         return_tensors="pt",
     ).input_ids
-    print("input_ids", input_ids)
+    print("input_ids:", input_ids)
 
     output_ids = dart.generate(
         input_ids,
@@ -80,10 +103,13 @@ def generate_prompt(aspect_ratio: str):
 
     generated = output_ids[len(input_ids) :]
     decoded = ", ".join([token for token in tokenizer.batch_decode(generated, skip_special_tokens=True) if token.strip() != ""])
-    print("decoded", decoded)
+    print("decoded:", decoded)
 
     return decoded
 
+def format_prompt(prompt: str, prompt_suffix: str):
+    return f"{prompt}, {prompt_suffix}"
+
 @spaces.GPU
 def generate_image(
     prompt: str,
@@ -93,7 +119,6 @@ def generate_image(
     height: int,
     guidance_scale: float,
     num_inference_steps: int,
-    progress=gr.Progress(track_tqdm=True),
 ):
     image = pipe(
         prompt=prompt,
@@ -108,6 +133,7 @@ def generate_image(
     return image
 
 def on_generate(
+    suffix: str,
     negative_prompt: str,
     seed,
     randomize_seed,
@@ -115,12 +141,15 @@ def on_generate(
     height,
     guidance_scale,
     num_inference_steps,
+    progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
 
-    prompt = generate_prompt("<|aspect_ratio:square|>")
+    ar = get_aspect_ratio(width, height)
+    prompt = generate_prompt(ar)
+    prompt = format_prompt(prompt, suffix)
     print(prompt)
 
     image = generate_image(
@@ -155,15 +184,21 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Generation details", open=False):
-            prompt_txt = gr.Textbox("Generated prompt", interactive=False)
+            prompt_txt = gr.Textbox(label="Generated prompt", interactive=False)
 
         with gr.Accordion("Advanced Settings", open=False):
+            prompt_suffix = gr.Text(
+                label="Prompt suffix",
+                max_lines=1,
+                visible=False,
+                value="masterpiece, best quality",
+            )
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
                 visible=False,
-                value=" worst quality, comic, multiple views, bad quality, low quality, lowres, displeasing, very displeasing, bad anatomy, bad hands, scan artifacts, monochrome, greyscale, signature, twitter username, jpeg artifacts, 2koma, 4koma, guro, extra digits, fewer digits",
+                value="worst quality, bad quality, low quality, lowres, displeasing, very displeasing, bad anatomy, bad hands, scan artifacts, signature, username, jpeg artifacts, guro, extra digits, fewer digits",
             )
 
             seed = gr.Slider(
@@ -214,6 +249,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click],
         fn=on_generate,
         inputs=[
+            prompt_suffix,
             negative_prompt,
             seed,
             randomize_seed,