Updates

app.py

@@ -20,11 +20,11 @@ if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶</p>"
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") != "0"
+CACHE_EXAMPLES = False #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
+PREVIEW_IMAGES = False
 
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -47,10 +47,12 @@ if torch.cuda.is_available():
         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):
+        def callback_prior(pipeline, step_index, t, callback_kwargs):
+            latents = callback_kwargs["latents"]
             output = previewer(latents)
             output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy())
-            return output
+            callback_kwargs["preview_output"] = output
+            return callback_kwargs
         callback_steps = 1
     else:
         previewer = None
@@ -62,6 +64,7 @@ else:
 
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    print("randomizing seed")
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
@@ -82,7 +85,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)
+    
+    #previewer.eval().requires_grad_(False).to(device).to(dtype)
     prior_pipeline.to(device)
     decoder_pipeline.to(device)
     
@@ -98,10 +102,9 @@ def generate(
         guidance_scale=prior_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
         generator=generator,
-        callback=callback_prior,
-        callback_steps=callback_steps
+        #callback_on_step_end=callback_prior,
+        #callback_on_step_end_tensor_inputs=['latents']
     )
-
     if PREVIEW_IMAGES:
         for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
             r = next(prior_output)
@@ -119,7 +122,7 @@ def generate(
         generator=generator,
         output_type="pil",
     ).images
-
+    print(decoder_output)
     #Save images
     for image in decoder_output:
         user_history.save_image(
@@ -137,15 +140,17 @@ def generate(
                 "num_images_per_prompt": num_images_per_prompt,
             },
         )
-
+    
     yield decoder_output[0]
 
 
 examples = [
-    "An astronaut riding a green horse",
-    "A mecha robot in a favela by Tarsila do Amaral",
-    "The sprirt of a Tamagotchi wandering in the city of Los Angeles",
-    "A delicious feijoada ramen dish"
+    "A futuristic cityscape at sunset",
+    "pair of shoes made of dried fruit skins, 3d render, bright colours, clean composition, beautiful artwork, logo",
+    "post-apocalyptic wasteland, the most delicate beautiful flower with green leaves growing from dust and rubble, vibrant colours, cinematic",
+    "Mixed media artwork, Emotional cyborg girl, Elegant dress, Skin lesions as a storytelling element, In the style of surrealist expressionism, muted color scheme, dreamlike atmosphere, abstract and distorted forms",
+    "rendering, side shot, falf-strange body with complex system equipment with hyper detail robot, gaze, sci-fi, gloomy environment, foggy with light shader, cyan and yellow illuminations, dramatic lighting, RTX shader, hyper detail texture with reflection, HDRI, cyborg, grunge, bolt, UHD",
+    "vintage Japanese postcard, in the style of Kentaro Miura, featuring a black cat holding a vinyl record in its paws, with vintage colors including light beige and red tones on a white background, very detailed artwork."
 ]
 
 with gr.Blocks() as demo:
@@ -277,4 +282,4 @@ with gr.Blocks(css="style.css") as demo_with_history:
         user_history.render()
 
 if __name__ == "__main__":
-    demo_with_history.queue(max_size=20).launch()
+    demo_with_history.queue(max_size=20).launch()

requirements.txt

@@ -1,5 +1,3 @@
-git+https://github.com/kashif/diffusers.git@diffusers-yield-callback
-https://gradio-builds.s3.amazonaws.com/aabb08191a7d94d2a1e9ff87b0d3c3987cd519c5/gradio-4.18.0-py3-none-any.whl
+diffusers
 accelerate
-safetensors
-transformers
+transformers