继续完善逻辑

app.py

@@ -138,9 +138,31 @@ def generate(
         pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, scheduler)
 
         upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
-        
 
-        return None
+        latents = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                output_type="latent",
+            ).images
+
+        upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+        images = upscaler_pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image=upscaled_latents,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                strength=upscaler_strength,
+                generator=generator,
+                output_type="pil",
+            ).images
+        return images[0]
+    
     except GenerationError as e:
         logger.warning(f"Generation validation error: {str(e)}")
         raise gr.Error(str(e))

utils.py

@@ -75,4 +75,21 @@ def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
         ),
         "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
     }
-    return scheduler_factory_map.get(name, lambda: None)()
+    return scheduler_factory_map.get(name, lambda: None)()
+
+def common_upscale(
+    samples: torch.Tensor,
+    width: int,
+    height: int,
+    upscale_method: str,
+) -> torch.Tensor:
+    return torch.nn.functional.interpolate(
+        samples, size=(height, width), mode=upscale_method
+    )
+
+def upscale(
+    samples: torch.Tensor, upscale_method: str, scale_by: float
+) -> torch.Tensor:
+    width = round(samples.shape[3] * scale_by)
+    height = round(samples.shape[2] * scale_by)
+    return common_upscale(samples, width, height, upscale_method)