fix: settings

README.md

@@ -1,6 +1,6 @@
 ---
-title: Random Illustrious Xl
-emoji: 🖼
+title: Random Illustrious XL
+emoji: 🎨🖌️
 colorFrom: purple
 colorTo: red
 sdk: gradio

app.py

@@ -36,13 +36,12 @@ IMAGE_MODEL_REPO_ID = os.environ.get(
 DART_V3_REPO_ID = os.environ.get("DART_V3_REPO_ID", None)
 assert DART_V3_REPO_ID is not None
 
-torch_dtype = torch.bfloat16
-
 dart = AutoModelForCausalLM.from_pretrained(
     DART_V3_REPO_ID,
-    torch_dtype=torch_dtype,
+    torch_dtype=torch.bfloat16,
     token=HF_TOKEN,
     use_cache=True,
+    device_map="cpu",
 )
 dart = dart.eval()
 dart = dart.requires_grad_(False)
@@ -50,14 +49,15 @@ dart = torch.compile(dart)
 tokenizer = AutoTokenizer.from_pretrained(DART_V3_REPO_ID)
 
 pipe = StableDiffusionXLPipeline.from_pretrained(
-    IMAGE_MODEL_REPO_ID, 
-    torch_dtype=torch_dtype,
+    IMAGE_MODEL_REPO_ID,
+    torch_dtype=torch.bfloat16,
     add_watermarker=False,
-    custom_pipeline="lpw_stable_diffusion_xl"
+    custom_pipeline="lpw_stable_diffusion_xl",
 )
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe.unet.set_attn_processor(AttnProcessor2_0())
-pipe = pipe.to(device)
+if device == "cuda":
+    pipe.enable_sequential_cpu_offload(gpu_id=0, device="cuda")
 
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -70,32 +70,35 @@ TEMPLATE = (
     "{aspect_ratio}"
     "<|length:medium|>"
     #
-    "<copyright>original</copyright>"
+    "<copyright></copyright>"
     #
     "<character></character>"
     #
-    "<general>"
+    "<general>{subject}"
 )
+QUALITY_TAGS = "masterpiece, best quality, very aesthetic, newest"
+NEGATIVE_PROMPT = "nsfw, (worst quality, bad quality:1.2), very displeasing, lowres, jaggy lines, 3d, watermark, signature, copyright, logo, blurry, ugly, poorly drawn, retro, scan, white outline"
+
 
 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|>"
+        return "<|aspect_ratio:ultra_tall|>"
+    elif ar <= 8 / 9:
+        return "<|aspect_ratio:tall|>"
     elif ar < 9 / 8:
         return "<|aspect_ratio:square|>"
     elif ar < math.sqrt(3):
-        return "<|aspect_ratio:tall|>"
+        return "<|aspect_ratio:wide|>"
     else:
-        return "<|aspect_ratio:ultra_tall|>"
+        return "<|aspect_ratio:ultra_wide|>"
 
 
 @torch.inference_mode
-def generate_prompt(aspect_ratio: str):
+def generate_prompt(subject: str, aspect_ratio: str):
     input_ids = tokenizer.encode_plus(
-        TEMPLATE.format(aspect_ratio=aspect_ratio),
+        TEMPLATE.format(aspect_ratio=aspect_ratio, subject=subject),
         return_tensors="pt",
     ).input_ids
     print("input_ids:", input_ids)
@@ -111,14 +114,22 @@ def generate_prompt(aspect_ratio: str):
     )[0]
 
     generated = output_ids[len(input_ids) :]
-    decoded = ", ".join([token for token in tokenizer.batch_decode(generated, skip_special_tokens=True) if token.strip() != ""])
+    decoded = ", ".join(
+        [
+            token
+            for token in tokenizer.batch_decode(generated, skip_special_tokens=True)
+            if token.strip() != ""
+        ]
+    )
     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,
@@ -141,7 +152,9 @@ def generate_image(
 
     return image
 
+
 def on_generate(
+    subject: str,
     suffix: str,
     negative_prompt: str,
     seed,
@@ -157,7 +170,8 @@ def on_generate(
     generator = torch.Generator().manual_seed(seed)
 
     ar = get_aspect_ratio(width, height)
-    prompt = generate_prompt(ar)
+    print("ar:", ar)
+    prompt = generate_prompt(subject, ar)
     prompt = format_prompt(prompt, suffix)
     print(prompt)
 
@@ -188,6 +202,11 @@ with gr.Blocks(css=css) as demo:
         """)
 
         with gr.Row():
+            subject_radio = gr.Dropdown(
+                label="Subject",
+                choices=["1girl", "2girls", "1boy", "no humans"],
+                value="1girl",
+            )
             run_button = gr.Button("Generate random", scale=0)
 
         result = gr.Image(label="Result", show_label=False)
@@ -199,13 +218,13 @@ with gr.Blocks(css=css) as demo:
             prompt_suffix = gr.Text(
                 label="Prompt suffix",
                 visible=True,
-                value="masterpiece, best quality, very aesthetic",
+                value=QUALITY_TAGS,
             )
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 placeholder="Enter a negative prompt",
                 visible=True,
-                value="(worst quality, bad quality, low quality:1.2), lowres, displeasing, very displeasing, bad anatomy, bad hands, extra digits, fewer digits, scan artifacts, signature, username, jpeg artifacts, retro, 2010s",
+                value=NEGATIVE_PROMPT,
             )
 
             seed = gr.Slider(
@@ -241,7 +260,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1.0,
                     maximum=10.0,
                     step=0.5,
-                    value=7,
+                    value=6.5,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -256,6 +275,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click],
         fn=on_generate,
         inputs=[
+            subject_radio,
             prompt_suffix,
             negative_prompt,
             seed,