Update app.py

app.py

@@ -360,43 +360,43 @@ def process(input_fg, input_bg, prompt, image_width, image_height, num_samples,
     ).images.to(vae.dtype) / vae.config.scaling_factor
 
     pixels = vae.decode(latents).sample
-    pixels = pytorch2numpy(pixels)
+    pixels = pytorch2numpy(pixels)  # Use default quant=True for first pass
     
-    if highres_scale > 1.0:
-        pixels = [resize_without_crop(
-            image=p,
-            target_width=int(round(image_width * highres_scale / 64.0) * 64),
-            target_height=int(round(image_height * highres_scale / 64.0) * 64))
-        for p in pixels]
-
-        pixels = numpy2pytorch(pixels).to(device=vae.device, dtype=vae.dtype)
-        latents = vae.encode(pixels).latent_dist.mode() * vae.config.scaling_factor
-        latents = latents.to(device=unet.device, dtype=unet.dtype)
-
-        image_height, image_width = latents.shape[2] * 8, latents.shape[3] * 8
-        fg = resize_and_center_crop(input_fg, image_width, image_height)
-        bg = resize_and_center_crop(input_bg, image_width, image_height)
-        concat_conds = numpy2pytorch([fg, bg]).to(device=vae.device, dtype=vae.dtype)
-        concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
-        concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)
-
-        latents = i2i_pipe(
-            image=latents,
-            strength=highres_denoise,
-            prompt_embeds=conds,
-            negative_prompt_embeds=unconds,
-            width=image_width,
-            height=image_height,
-            num_inference_steps=int(round(steps / highres_denoise)),
-            num_images_per_prompt=num_samples,
-            generator=rng,
-            output_type='latent',
-            guidance_scale=cfg,
-            cross_attention_kwargs={'concat_conds': concat_conds},
-        ).images.to(vae.dtype) / vae.config.scaling_factor
-
-        pixels = vae.decode(latents).sample
-        pixels = pytorch2numpy(pixels, quant=False)
+    # Always perform highres processing like the original code
+    pixels = [resize_without_crop(
+        image=p,
+        target_width=int(round(image_width * highres_scale / 64.0) * 64),
+        target_height=int(round(image_height * highres_scale / 64.0) * 64))
+    for p in pixels]
+
+    pixels = numpy2pytorch(pixels).to(device=vae.device, dtype=vae.dtype)
+    latents = vae.encode(pixels).latent_dist.mode() * vae.config.scaling_factor
+    latents = latents.to(device=unet.device, dtype=unet.dtype)
+
+    image_height, image_width = latents.shape[2] * 8, latents.shape[3] * 8
+    fg = resize_and_center_crop(input_fg, image_width, image_height)
+    bg = resize_and_center_crop(input_bg, image_width, image_height)
+    concat_conds = numpy2pytorch([fg, bg]).to(device=vae.device, dtype=vae.dtype)
+    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
+    concat_conds = torch.cat([c[None, ...] for c in concat_conds], dim=1)
+
+    latents = i2i_pipe(
+        image=latents,
+        strength=highres_denoise,
+        prompt_embeds=conds,
+        negative_prompt_embeds=unconds,
+        width=image_width,
+        height=image_height,
+        num_inference_steps=int(round(steps / highres_denoise)),
+        num_images_per_prompt=num_samples,
+        generator=rng,
+        output_type='latent',
+        guidance_scale=cfg,
+        cross_attention_kwargs={'concat_conds': concat_conds},
+    ).images.to(vae.dtype) / vae.config.scaling_factor
+
+    pixels = vae.decode(latents).sample
+    pixels = pytorch2numpy(pixels, quant=False)  # Return 0-1 range floats for final result
 
     return pixels, [fg, bg]