Update app.py

app.py

@@ -1,6 +1,7 @@
 # import os
 import spaces
 
+import time
 import gradio as gr
 import torch
 from PIL import Image
@@ -62,7 +63,7 @@ class HFEmbedder(nn.Module):
     
 
 device = "cuda"
-t5 = HFEmbedder("google/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
+t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
 clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
 ae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
 # quantize(t5, weights=qfloat8)
@@ -760,14 +761,15 @@ def generate_image(
     if seed == 0:
         seed = int(random.random() * 1000000)
         
-    assert torch.cuda.is_available()
-    torch_device = torch.device("cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch_device = torch.device(device)
     
     global model, model_zero_init
     if not model_zero_init:
         model = model.to(torch_device)
         model_zero_init = True
-
+        
+    t = time.perf_counter()
     if do_img2img and init_image is not None:
         init_image = get_image(init_image)
         if resize_img:
@@ -792,9 +794,8 @@ def generate_image(
         timesteps = timesteps[t_idx:]
         x = t * x + (1.0 - t) * init_image.to(x.dtype)
 
-    with torch_device:
-        inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
-        x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
+    inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
+    x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
     
     # with profile(activities=[ProfilerActivity.CPU],record_shapes=True,profile_memory=True) as prof:
     # print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=20))
@@ -807,7 +808,8 @@ def generate_image(
     x = x.clamp(-1, 1)
     x = rearrange(x[0], "c h w -> h w c")
     img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
-
+    seed += f"\nGenerated in {time.perf_counter() - t:.2f}s"
+    
     return img, seed
 
 def create_demo():