Update app.py

app.py

@@ -72,12 +72,10 @@ pipe_controlnet = FluxControlNetPipeline(
     tokenizer_2=pipe.tokenizer_2,
     transformer=pipe.transformer,  # unet 대신 transformer 사용
     controlnet=controlnet,
-    scheduler=pipe.scheduler
+    scheduler=pipe.scheduler,
+    torch_dtype=dtype
 ).to(device)
 
-
-
-
 MAX_SEED = 2**32 - 1
 MAX_PIXEL_BUDGET = 1024 * 1024
 
@@ -310,8 +308,8 @@ def remove_custom_lora(selected_indices, current_loras):
 @spaces.GPU(duration=75)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
     print("Generating image...")
-    pipe.to("cuda")
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    pipe.to(device)
+    generator = torch.Generator(device=device).manual_seed(seed)
     with calculateDuration("Generating image"):
         # Generate image
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -329,8 +327,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
 
 @spaces.GPU(duration=75)
 def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed):
-    pipe_i2i.to("cuda")
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    pipe_i2i.to(device)
+    generator = torch.Generator(device=device).manual_seed(seed)
     image_input = load_image(image_input_path)
     final_image = pipe_i2i(
         prompt=prompt_mash,
@@ -414,7 +412,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     # Generate image
     if image_input is not None:
         final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
-        yield final_image, seed, gr.update(visible=False)
+        return final_image, seed, gr.update(visible=False)
     else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
         # Consume the generator to get the final image
@@ -518,8 +516,8 @@ def process_input(input_image, upscale_factor, **kwargs):
 
     return input_image.resize((w, h)), w_original, h_original, was_resized
 
-@spaces.GPU#(duration=42)
-def infer(
+@spaces.GPU(duration=75)
+def infer_upscale(
     seed,
     randomize_seed,
     input_image,
@@ -561,8 +559,8 @@ def infer(
     # resize to target desired size
     image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
     image.save("output.jpg")
-    # convert to numpy
-    return [true_input_image, image, seed]
+    # convert to PIL Image
+    return [true_input_image, image]
 
 css = '''
 #gen_btn{height: 100%}
@@ -592,6 +590,7 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
 
     loras_state = gr.State(loras)
     selected_indices = gr.State([])
+    
     with gr.Tab("Generate"):
         with gr.Row():
             with gr.Column(scale=3):
@@ -660,44 +659,51 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
                         randomize_seed = gr.Checkbox(True, label="Randomize seed")
                         seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
 
+        # 이벤트 핸들러 설정
+        generate_button.click(
+            fn=run_lora,
+            inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
+            outputs=[result, seed, progress_bar]
+        )
+        prompt.submit(
+            fn=run_lora,
+            inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
+            outputs=[result, seed, progress_bar]
+        )
         gallery.select(
-            update_selection,
+            fn=update_selection,
             inputs=[selected_indices, loras_state, width, height],
-            outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2])
+            outputs=[prompt, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, width, height, lora_image_1, lora_image_2]
+        )
         remove_button_1.click(
-            remove_lora_1,
+            fn=remove_lora_1,
             inputs=[selected_indices, loras_state],
             outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
         )
         remove_button_2.click(
-            remove_lora_2,
+            fn=remove_lora_2,
             inputs=[selected_indices, loras_state],
             outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
         )
         randomize_button.click(
-            randomize_loras,
+            fn=randomize_loras,
             inputs=[selected_indices, loras_state],
             outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, prompt]
         )
         add_custom_lora_button.click(
-            add_custom_lora,
+            fn=add_custom_lora,
             inputs=[custom_lora, selected_indices, loras_state],
             outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
         )
         remove_custom_lora_button.click(
-            remove_custom_lora,
+            fn=remove_custom_lora,
             inputs=[selected_indices, loras_state],
             outputs=[loras_state, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
         )
-        gr.on(
-            triggers=[generate_button.click, prompt.submit],
-            fn=run_lora,
-            inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, loras_state],
-            outputs=[result, seed, progress_bar]
-        ).then(  # Update the history gallery
-            fn=lambda x, history: update_history(x, history),
+        history_gallery.load(
+            fn=update_history,
             inputs=[result, history_gallery],
-            outputs=history_gallery,
+            outputs=history_gallery
         )
 
     with gr.Tab("Upscale"):
@@ -737,8 +743,9 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         with gr.Row():
             upscale_button = gr.Button("Upscale", variant="primary")
 
+        # 업스케일 버튼 이벤트 핸들러
         upscale_button.click(
-            infer,
+            fn=infer_upscale,
             inputs=[
                 seed_upscale,
                 randomize_seed_upscale,
@@ -752,3 +759,4 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
 
 app.queue()
 app.launch()
+