Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import gc
 import random
 from typing import Callable, Dict, Optional, Tuple
 
@@ -9,7 +10,7 @@ import spaces
 import torch
 
 from transformers import CLIPTextModel
-from diffusers import AutoencoderKL, StableDiffusionXLPipeline, DDIMScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
+from diffusers import AutoencoderKL, StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, DDIMScheduler, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler
 
 MODEL = "eienmojiki/Starry-XL-v5.2"
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -93,6 +94,28 @@ def load_pipeline(model_name):
     pipe.to(device)
     return pipe
 
+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)
+
+def free_memory() -> None:
+    torch.cuda.empty_cache()
+    gc.collect()
+
 @spaces.GPU(enable_queue=False)
 def generate(
     prompt: str,
@@ -118,9 +141,11 @@ def generate(
     
     pipe.to(device)
 
+    upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
+
     try:
 
-        img = pipe(
+        latents = pipe(
             prompt = prompt,
             negative_prompt = negative_prompt,
             width = width,
@@ -129,6 +154,19 @@ def generate(
             num_inference_steps = num_inference_steps,
             generator = generator,
             num_images_per_prompt=1,
+            output_type="latents",
+        ).images
+
+        upscaled_latents = upscale(latents, "nearest-exact", 2.0)
+
+        img = upscaler_pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=upscaled_latents,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            strength=0.55,
+            generator=generator,
             output_type="pil",
         ).images[0]
 
@@ -136,6 +174,9 @@ def generate(
 
     except Exception as e:
         print(f"An error occurred: {e}")
+    finally:
+        del upscaler_pipe
+        free_memory()
 
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)