Update app.py

app.py

@@ -72,7 +72,7 @@ def get_seed(seed):
         return random.randint(0, MAX_SEED)
 
 @spaces.GPU(duration=60)
-def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATIVE_INPUT, model=DEFAULT_MODEL, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, steps=1, guidance=0, number=1, seed=None):
+def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATIVE_INPUT, model=DEFAULT_MODEL, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, steps=1, guidance=0, number=1, seed=None, height_buffer=DEFAULT_HEIGHT, width_buffer=DEFAULT_WIDTH):
 
     repo = repo_customs[model or "Default"]
     filter_input = filter_input or ""
@@ -133,15 +133,12 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
 
     print(nsfw_prediction)
 
-    image = images[0].convert("RGBA").resize((width, height))
-    buffer_data = []
-
-    for y in range(height):
-        for x in range(width):
-            r, g, b, a = image.getpixel((x, y))
-            buffer_data.extend([r, g, b, a])
-
-    buffer_json = json.dumps(buffer_data)
+    buffer_image = images[0].convert("RGBA").resize((width_buffer, height_buffer))
+    
+    image_array = np.array(buffer_image)
+    pixel_data = image_array.flatten().tolist()
+    
+    buffer_json = json.dumps(pixel_data)
 
     return image_paths, {item['label']: round(item['score'], 3) for item in nsfw_prediction}, buffer_json
 
@@ -164,6 +161,8 @@ with gr.Blocks(css=css) as main:
         guidance = gr.Slider(minimum=0, maximum=100, step=0.1, value=5, label = "Guidance")
         number = gr.Slider(minimum=1, maximum=4, step=1, value=1, label="Number")
         seed = gr.Textbox(lines=1, value="", label="Seed (Blank for random)")
+        height_buffer = gr.Slider(minimum=1, maximum=2160, step=1, value=DEFAULT_HEIGHT, label="Buffer Height")
+        width_buffer = gr.Slider(minimum=1, maximum=2160, step=1, value=DEFAULT_WIDTH, label="uffer Width")
         submit = gr.Button("▶")
         maintain = gr.Button("☁️")
 
@@ -172,7 +171,7 @@ with gr.Blocks(css=css) as main:
         output_2 = gr.Label()
         output_3 = gr.Textbox(lines=1, value="", label="Buffer")
             
-    submit.click(generate, inputs=[input, filter_input, negative_input, model, height, width, steps, guidance, number, seed], outputs=[output, output_2], queue=False)
+    submit.click(generate, inputs=[input, filter_input, negative_input, model, height, width, steps, guidance, number, seed, height_buffer, width_buffer], outputs=[output, output_2], queue=False)
     maintain.click(cloud, inputs=[], outputs=[], queue=False)
 
 main.launch(show_api=True)