Update app.py

app.py

@@ -158,6 +158,7 @@ validation_pipeline._init_tiled_vae(encoder_tile_size=1024,
 weight_dtype = torch.float16
 device = "cuda"
 
+
 # Move text_encode and vae to gpu and cast to weight_dtype
 text_encoder.to(device, dtype=weight_dtype)
 vae.to(device, dtype=weight_dtype)
@@ -173,17 +174,27 @@ tag_model.eval()
 tag_model.to(device, dtype=weight_dtype)
 
 def preprocess_image(input_image: Image.Image) -> Image.Image:
-    img = input_image.copy()
-    img.thumbnail((1024, 1024), Image.Resampling.BILINEAR)
-    return img
 
+    ori_width, ori_height = input_image.size
+
+    duration = 60
+
+    if ori_width > 384 or ori_height > 384:
+        duration = 75
+    elif ori_width > 512 or ori_height > 512:
+        duration = 90
+    elif ori_width > 1024 or ori_height > 1024:
+        duration = 120
+
+    return duration
 
 @spaces.GPU()
 def preprocess_n_magnify(input_image: Image.Image):
-    preprocessed_img = preprocess_image(input_image)
-    preprocessed_img, magnified_img = magnify(preprocessed_img)
+    duration = preprocess_image(input_image)
+    magnified_img = magnify(input_image, duration)    
+
+    return duration, (input_image, magnified_img)
 
-    return preprocessed_img, (preprocessed_img, magnified_img)
 
 def get_duration(
     input_image: Image.Image,
@@ -201,12 +212,7 @@ def get_duration(
     progress,
     ):
     
-    if steps > 4 and duration_seconds > 2:
-        return 90
-    elif steps > 4 or duration_seconds > 2:
-        return 75
-    else:
-        return 60
+    return duration_seconds
         
     
 @spaces.GPU(duration=get_duration)
@@ -227,7 +233,6 @@ def magnify(
     ):
 
     
-    
     process_size = 512
     resize_preproc = transforms.Compose([
         transforms.Resize(process_size, interpolation=transforms.InterpolationMode.BILINEAR),
@@ -319,7 +324,7 @@ with gr.Blocks(css=css, theme=theme) as demo:
             with gr.Column():
                 input_image = gr.Image(type="pil", height=256)
                 run_button = gr.Button("🔎 Magnify 4x", variant="primary")
-                preprocessed_image = gr.Image(label="preprocessed image", type="pil", interactive=False, height=256)
+                duration_time = gr.Text(label="duration time", value=60)
                 with gr.Accordion("Options", visible=False):
                     user_prompt = gr.Textbox(label="User Prompt", value="")
                     positive_prompt = gr.Textbox(label="Positive Prompt", value="clean, high-resolution, 8k, best quality, masterpiece")
@@ -356,14 +361,14 @@ with gr.Blocks(css=css, theme=theme) as demo:
                     inputs=[
                         input_image,
                     ],
-                    outputs=[preprocessed_image, result_gallery],
+                    outputs=[duration_time, result_gallery],
                     fn=preprocess_n_magnify,
                     cache_examples=True,
                 )
         inputs = [
             input_image,
         ]
-        run_button.click(fn=magnify, inputs=preprocessed_image, outputs=[result_gallery])
-        input_image.upload(fn=preprocess_image,inputs=input_image, outputs=[preprocessed_image])
+        run_button.click(fn=magnify, inputs=[input_image,duration_time], outputs=[result_gallery])
+        input_image.upload(fn=preprocess_image,inputs=input_image, outputs=duration_time)
 
 demo.launch(share=True)