Update app.py

app.py

@@ -3,16 +3,20 @@ import torch
 import numpy as np
 import modin.pandas as pd
 from PIL import Image
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionLatentUpscalePipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = DiffusionPipeline.from_pretrained("prompthero/openjourney-v4", torch_dtype=torch.float16, safety_checker=None)
+upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
+upscaler = upscaler.to(device)
 pipe = pipe.to(device)
 
 def genie (prompt, scale, steps, Seed):
-     generator = torch.Generator(device=device).manual_seed(Seed)
-     images = pipe(prompt, num_inference_steps=steps, guidance_scale=scale, generator=generator).images[0]
-     return images
+    generator = torch.Generator(device=device).manual_seed(Seed)
+     #images = pipe(prompt, num_inference_steps=steps, guidance_scale=scale, generator=generator).images[0]
+    low_res_latents = pipeline(prompt, num_inference_steps=steps, guidance_scale=scale, generator=generator, output_type="latent").images
+    upscaled_image = upscaler(prompt=prompt, image=low_res_latents, num_inference_steps=20, guidance_scale=0, generator=generator).images[0]
+    return upscaled_image
     
 gr.Interface(fn=genie, inputs=[gr.Textbox(label='What you want the AI to generate. 77 Token Limit.'), 
                                gr.Slider(1, maximum=25, value=10, step=.25, label='Prompt Guidance Scale:', interactive=True),