stable1

README.md

@@ -4,7 +4,7 @@ emoji: 🖼️🖌️
 colorFrom: yellow
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.27.0
+sdk_version: 4.31.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,9 +1,10 @@
-
+import gc
 import json
 import webcolors
 import spaces
 import gradio as gr
 import os.path as osp
+from copy import deepcopy
 from PIL import Image, ImageDraw, ImageFont
 
 import torch
@@ -64,6 +65,10 @@ font = ImageFont.truetype("assets/Arial.ttf", 20)
 
 device = "cuda"
 
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+
 def import_model_class_from_model_name_or_path(
     pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder",
 ):
@@ -215,6 +220,18 @@ pipeline.scheduler = DPMSolverMultistepScheduler.from_pretrained(
 
 prompt_format = PromptFormat()
 
+# move to gpu
+if config.pretrained_vae_model_name_or_path is None:
+    vae = vae.to(device, dtype=torch.float32)
+else:
+    vae = vae.to(device, dtype=inference_dtype)
+text_encoder_one = text_encoder_one.to(device, dtype=inference_dtype)
+text_encoder_two = text_encoder_two.to(device, dtype=inference_dtype)
+byt5_model = byt5_model.to(device)
+unet = unet.to(device, dtype=inference_dtype)
+pipeline = pipeline.to(device)
+
+
 def get_pixels(
     box_sketch_template,
     evt: gr.SelectData
@@ -318,8 +335,6 @@ def exe_undo(
     return box_sketch_template
 
 def process_box():
-    global stack
-    global state
 
     visibilities = []
     for _ in range(MAX_TEXT_BOX + 1):
@@ -330,31 +345,19 @@ def process_box():
     # return [gr.update(visible=True), binary_matrixes, *visibilities, *colors]
     return [gr.update(visible=True), *visibilities]
 
-@spaces.GPU
+@torch.inference_mode()
+@spaces.GPU(enable_queue=True)
 def generate_image(bg_prompt, bg_class, bg_tags, seed, *conditions):
-    print(conditions)
-    
-    # 0 load model to cuda
-    global pipeline
-    if config.pretrained_vae_model_name_or_path is None:
-        vae.to(device, dtype=torch.float32)
-    else:
-        vae.to(device, dtype=inference_dtype)
-    text_encoder_one.to(device, dtype=inference_dtype)
-    text_encoder_two.to(device, dtype=inference_dtype)
-    byt5_model.to(device)
-    unet.to(device, dtype=inference_dtype)
-    pipeline = pipeline.to(device)
 
+    stack_cp = deepcopy(stack)
+    print(f"conditions: {conditions}")
+    
     # 1. parse input
-    global state
-    global stack
-
     prompts = []
     colors = []
     font_type = []
     bboxes = []
-    num_boxes = len(stack) if len(stack[-1]) == 4 else len(stack) - 1
+    num_boxes = len(stack_cp) if len(stack_cp[-1]) == 4 else len(stack_cp) - 1
     for i in range(num_boxes):
         prompts.append(conditions[i])
         colors.append(conditions[i + MAX_TEXT_BOX])
@@ -373,10 +376,10 @@ def generate_image(bg_prompt, bg_class, bg_tags, seed, *conditions):
             raise gr.Error(f"Invalid style for text box {i + 1} !")
         bboxes.append(
             [
-                stack[i][0] / 1024,
-                stack[i][1] / 1024,
-                (stack[i][2] - stack[i][0]) / 1024,
-                (stack[i][3] - stack[i][1]) / 1024,
+                stack_cp[i][0] / 1024,
+                stack_cp[i][1] / 1024,
+                (stack_cp[i][2] - stack_cp[i][0]) / 1024,
+                (stack_cp[i][3] - stack_cp[i][1]) / 1024,
             ]
         )
         styles.append(
@@ -393,14 +396,11 @@ def generate_image(bg_prompt, bg_class, bg_tags, seed, *conditions):
         bg_prompt += " Tags: " + bg_tags
     text_prompt = prompt_format.format_prompt(prompts, styles)
 
-    print(bg_prompt)
-    print(text_prompt)
+    print(f"bg_prompt: {bg_prompt}")
+    print(f"text_prompt: {text_prompt}")
 
     # 4. inference
-    if seed == -1:
-        generator = torch.Generator(device=device)
-    else:
-        generator = torch.Generator(device=device).manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(int(seed))
     with torch.cuda.amp.autocast():
         image = pipeline(
             prompt=bg_prompt,
@@ -411,6 +411,9 @@ def generate_image(bg_prompt, bg_class, bg_tags, seed, *conditions):
             generator=generator,
             text_attn_mask=None,
         ).images[0]
+
+    flush()
+
     return image
 
 def process_example(bg_prompt, bg_class, bg_tags, color_str, style_str, text_str, box_str, seed):
@@ -534,10 +537,10 @@ def main():
                                     choices=font_idx_list,
                                 ))
 
-                    seed_ = gr.Slider(label="Seed", minimum=-1, maximum=999999999, value=-1, step=1)
-                    button_generate = gr.Button("(2) I've finished my texts, colors and styles, generate!", elem_id="main_button", interactive=True)
+                    seed_ = gr.Slider(label="Seed", minimum=0, maximum=2147483647, value=42, step=1)
+                    button_generate = gr.Button("(2) I've finished my texts, colors and styles, generate!", elem_id="main_button", interactive=True, variant='primary')
 
-                button_layout.click(process_box, inputs=[], outputs=[post_box, *color_row], queue=False)
+                button_layout.click(process_box, inputs=[], outputs=[post_box, *color_row])
 
             with gr.Column():
                 output_image = gr.Image(label="Output Image", interactive=False)
@@ -570,7 +573,7 @@ def main():
                     'LilitaOne, Sensei-Medium, Sensei-Medium, LilitaOne, LilitaOne, LilitaOne',
                     "RSVP to +123-456-7890**********Olivia Wilson**********Baby Shower**********Please Join Us For a**********In Honoring**********23 November, 2021 | 03:00 PM Fauget Hotels",
                     '[0.07112462006079028, 0.6462006079027356, 0.3373860182370821, 0.026747720364741642]; [0.07051671732522796, 0.38662613981762917, 0.37264437689969604, 0.059574468085106386]; [0.07234042553191489, 0.15623100303951368, 0.6547112462006079, 0.12401215805471125]; [0.0662613981762918, 0.06747720364741641, 0.3981762917933131, 0.035866261398176294]; [0.07051671732522796, 0.31550151975683893, 0.22006079027355624, 0.03951367781155015]; [0.06990881458966565, 0.48328267477203646, 0.39878419452887537, 0.1094224924012158]',
-                    0,
+                    1,
                 ],
                 [
                     'The image features a white background with a variety of colorful flowers and decorations. There are several pink flowers scattered throughout the scene, with some positioned closer to the top and others near the bottom. A blue flower can also be seen in the middle of the image. The overall composition creates a visually appealing and vibrant display.',
@@ -605,6 +608,7 @@ def main():
             ],
             outputs=[post_box, box_sketch_template, seed_, *color_row, *colors, *styles, *prompts],
             fn=process_example,
+            cache_examples=False,
             run_on_click=True,
             label='Examples',
         )

requirements.txt

@@ -7,4 +7,4 @@ torchvision==0.17.0
 deepspeed
 peft
 webcolors
-gradio
+gradio==4.31.1