Update app.py

app.py

@@ -8,7 +8,6 @@ from typing import List
 from diffusers.utils import numpy_to_pil
 from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
-import spaces 
 from previewer.modules import Previewer
 import user_history
 
@@ -24,14 +23,12 @@ CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") != "0
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
 USE_TORCH_COMPILE = False
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
-PREVIEW_IMAGES = True
 
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-    prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype)#.to(device)
-    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype)#.to(device) 
-
+    prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype).to(device)
+    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype).to(device) 
     if ENABLE_CPU_OFFLOAD:
         prior_pipeline.enable_model_cpu_offload()
         decoder_pipeline.enable_model_cpu_offload()
@@ -43,19 +40,6 @@ if torch.cuda.is_available():
         prior_pipeline.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
         decoder_pipeline.decoder = torch.compile(decoder_pipeline.decoder, mode="max-autotune", fullgraph=True)
     
-    if PREVIEW_IMAGES:
-        previewer = Previewer()
-        previewer_state_dict = torch.load("previewer/previewer_v1_100k.pt", map_location=torch.device('cpu'))["state_dict"]
-        previewer.load_state_dict(previewer_state_dict)
-        def callback_prior(i, t, latents):
-            output = previewer(latents)
-            output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy())
-            return output
-        callback_steps = 1
-    else:
-        previewer = None
-        callback_prior = None
-        callback_steps = None
 else:
     prior_pipeline = None
     decoder_pipeline = None
@@ -66,7 +50,6 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -82,12 +65,8 @@ def generate(
     num_images_per_prompt: int = 2,
     profile: gr.OAuthProfile | None = None,
 ) -> PIL.Image.Image:
-    previewer.eval().requires_grad_(False).to(device).to(dtype)
-    prior_pipeline.to(device)
-    decoder_pipeline.to(device)
     
     generator = torch.Generator().manual_seed(seed)
-    print("prior_num_inference_steps: ", prior_num_inference_steps)
     prior_output = prior_pipeline(
         prompt=prompt,
         height=height,
@@ -98,17 +77,8 @@ def generate(
         guidance_scale=prior_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
         generator=generator,
-        callback=callback_prior,
-        callback_steps=callback_steps
     )
 
-    if PREVIEW_IMAGES:
-        for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
-            r = next(prior_output)
-            if isinstance(r, list):
-                yield r[0]
-        prior_output = r
-
     decoder_output = decoder_pipeline(
         image_embeddings=prior_output.image_embeddings,
         prompt=prompt,
@@ -120,25 +90,7 @@ def generate(
         output_type="pil",
     ).images
 
-    #Save images
-    for image in decoder_output:
-        user_history.save_image(
-            profile=profile,
-            image=image,
-            label=prompt,
-            metadata={
-                "negative_prompt": negative_prompt,
-                "seed": seed,
-                "width": width,
-                "height": height,
-                "prior_guidance_scale": prior_guidance_scale,
-                "decoder_num_inference_steps": decoder_num_inference_steps,
-                "decoder_guidance_scale": decoder_guidance_scale,
-                "num_images_per_prompt": num_images_per_prompt,
-            },
-        )
-
-    yield decoder_output[0]
+    return decoder_output[0]
 
 
 examples = [
@@ -270,11 +222,8 @@ with gr.Blocks() as demo:
         api_name="run",
     )
     
-with gr.Blocks(css="style.css") as demo_with_history:
-    with gr.Tab("App"):
-        demo.render()
-    with gr.Tab("Past generations"):
-        user_history.render()
-
+with gr.Blocks(css="style.css") as local_demo:
+    demo.render()
+    
 if __name__ == "__main__":
-    demo_with_history.queue(max_size=20).launch()
+    local_demo.queue(max_size=20).launch()