Update app.py

app.py

@@ -144,14 +144,23 @@ def process_image(
             generator.manual_seed(seed)
         
         # Process input image
-        input_pil = Image.fromarray(input_image)
-        width, height = input_pil.size
+        validation_image = Image.fromarray(input_image)
+        ori_width, ori_height = validation_image.size
         
-        # Resize image
-        target_width = width * upscale_factor
-        target_height = height * upscale_factor
-        target_width = target_width - (target_width % 8)
-        target_height = target_height - (target_height % 8)
+        # Resize logic from original script
+        resize_flag = False
+        rscale = upscale_factor
+        process_size = 512  # Same as args.process_size in original
+        
+        if ori_width < process_size//rscale or ori_height < process_size//rscale:
+            scale = (process_size//rscale)/min(ori_width, ori_height)
+            tmp_image = validation_image.resize((round(scale*ori_width), round(scale*ori_height)))
+            validation_image = tmp_image
+            resize_flag = True
+
+        validation_image = validation_image.resize((validation_image.size[0]*rscale, validation_image.size[1]*rscale))
+        validation_image = validation_image.resize((validation_image.size[0]//8*8, validation_image.size[1]//8*8))
+        width, height = validation_image.size
         
         # Move pipeline to GPU for processing
         pipeline.to(accelerator.device)
@@ -164,12 +173,12 @@ def process_image(
                 False,  # tile_diffusion
                 None,   # tile_diffusion_size
                 None,   # tile_diffusion_stride
-                prompt, # validation_prompt / added_prompt
-                input_pil,  # validation_image
+                prompt,
+                validation_image,
                 num_inference_steps=num_inference_steps,
                 generator=generator,
-                height=target_height,
-                width=target_width,
+                height=height,
+                width=width,
                 guidance_scale=guidance_scale,
                 negative_prompt=negative_prompt,
                 conditioning_scale=conditioning_scale,
@@ -178,18 +187,21 @@ def process_image(
                 use_vae_encode_condition=False
             )
         
-        generated_image = output.images[0]
-        
+        image = output.images[0]
+
         # Apply color fixing if specified
         if color_fix_method != "none":
             fix_func = wavelet_color_fix if color_fix_method == "wavelet" else adain_color_fix
-            generated_image = fix_func(generated_image, input_pil)
+            image = fix_func(image, validation_image)
+            
+        if resize_flag:
+            image = image.resize((ori_width*rscale, ori_height*rscale))
         
         # Move pipeline back to CPU
         pipeline.to("cpu")
         torch.cuda.empty_cache()
         
-        return generated_image
+        return image
     
     except Exception as e:
         print(f"Error processing image: {str(e)}")