Update app.py

app.py

@@ -7,7 +7,7 @@ from fastapi.staticfiles import StaticFiles
 
 css_code = os.getenv("DazDinGo_CSS")
 
-# Buat folder /tmp jika belum ada
+# Create /tmp directory if it doesn't exist
 os.makedirs("/tmp/flux_images", exist_ok=True)
 
 def load_fn(models):
@@ -23,9 +23,6 @@ def load_fn(models):
 
 load_fn(models)
 
-num_models = len(models)
-default_models = models[:num_models]
-
 def save_image_to_tmp(image):
     if image is None:
         return None
@@ -41,58 +38,81 @@ def gen_fn(model_str, prompt):
     klir = '| ultra detail, ultra elaboration, ultra quality, perfect'
     result = models_load[model_str](f'{prompt} {klir} {noise}')
     
-    # Simpan gambar ke /tmp dan dapatkan URL
     if isinstance(result, list):
         return [save_image_to_tmp(img) for img in result]
-    else:
-        return save_image_to_tmp(result)
+    return save_image_to_tmp(result)
 
-def make_me():
-    with gr.Row():
-        with gr.Column(scale=4):
-            txt_input = gr.Textbox(label='Your prompt:', lines=4, container=False, 
-                                 elem_id="custom_textbox", placeholder="Prompt", height=250)
-        
-        with gr.Column(scale=1):
-            gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
-            stop_button = gr.Button('Stop', variant='secondary', interactive=False, 
-                                   elem_id="custom_stop_button")
+def create_interface():
+    with gr.Blocks(css=css_code) as demo:
+        with gr.Row():
+            with gr.Column(scale=4):
+                txt_input = gr.Textbox(
+                    label='Your prompt:', 
+                    lines=4, 
+                    container=False,
+                    elem_id="custom_textbox", 
+                    placeholder="Prompt"
+                )
             
-            def on_generate_click():
-                return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
-            
-            def on_stop_click():
-                return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+            with gr.Column(scale=1):
+                gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
+                stop_button = gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                def toggle_buttons():
+                    return [
+                        gr.Button(interactive=False),
+                        gr.Button(interactive=True)
+                    ]
+                
+                gen_button.click(
+                    fn=toggle_buttons,
+                    outputs=[gen_button, stop_button]
+                )
+                
+                stop_button.click(
+                    fn=lambda: [gr.Button(interactive=True), gr.Button(interactive=False)],
+                    outputs=[gen_button, stop_button]
+                )
+        
+        with gr.Row():
+            output = gr.Gallery(
+                label="Generated Images",
+                columns=3,
+                height="auto",
+                object_fit="contain",
+                show_download_button=True,
+                elem_id="custom_gallery"
+            )
             
-            gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
-            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
-    
-    with gr.Row():
-        # Ubah output menjadi Gallery untuk menampilkan gambar yang bisa di-download
-        output = gr.Gallery(label="Generated Images", 
-                          columns=3, 
-                          height="auto",
-                          object_fit="contain",
-                          show_download_button=True,
-                          elem_id="custom_gallery")
+            model_selector = gr.Dropdown(
+                models,
+                label="Select LoRA Model",
+                value=models[0] if models else None,
+                multiselect=False,
+                interactive=True,
+                elem_id="custom_dropdown"
+            )
         
-        model_selector = gr.Dropdown(models, 
-                                   label="Select LoRA Model", 
-                                   value=models[0] if models else None,
-                                   multiselect=False,
-                                   interactive=True,
-                                   elem_id="custom_dropdown")
-    
-    with gr.Row():
-        gen_button.click(gen_fn, [model_selector, txt_input], output)
-        stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_button.click])
+        # Store the generation event to be able to cancel it
+        gen_event = gen_button.click(
+            fn=gen_fn,
+            inputs=[model_selector, txt_input],
+            outputs=output
+        )
+        
+        stop_button.click(
+            fn=None,
+            cancels=[gen_event],
+            outputs=[gen_button, stop_button]
+        )
+        
+        return demo
 
-with gr.Blocks(css=css_code) as demo: 
-    make_me()
+demo = create_interface()
 
-# Mount static files after demo is defined
+# Configure static files serving
 app = demo.app
 app.mount("/file", StaticFiles(directory="/tmp"), name="static")
 
-demo.queue(concurrency_count=50)
-demo.launch()
+if __name__ == "__main__":
+    demo.queue(concurrency_count=50).launch()