Update app.py

app.py

@@ -77,7 +77,7 @@ class ModelWrapper:
     def _get_time():
         return time.time()
         
-    @spaces.GPU()
+    @spaces.GPU(duration=100)
     def sample(self, noise, unet_added_conditions, prompt_embed, fast_vae_decode):
         alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
 
@@ -89,13 +89,15 @@ class ModelWrapper:
             step_interval = 250 
         else:
             raise NotImplementedError()
-        
+
+        noise = noise.to(device="cuda", dtype=torch.float16)
+        print(f'noise: {noise.dtype}')
         DTYPE = prompt_embed.dtype
         print(f'prompt_embed: {DTYPE}')
         
         for constant in all_timesteps:
             current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
-            current_timesteps = current_timesteps.to(torch.float16)
+            #current_timesteps = current_timesteps.to(torch.float16)
             print(f'current_timestpes: {current_timesteps.dtype}')
             eval_images = self.model(noise, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions).sample
             print(type(eval_images))
@@ -124,7 +126,7 @@ class ModelWrapper:
 
         add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)
 
-        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda", dtype=torch.float16)
+        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator)
         
         prompt_inputs = self._encode_prompt(prompt)
         
@@ -143,7 +145,7 @@ class ModelWrapper:
         }
 
 
-        print(f'noise: {noise.dtype}')
+        
         print(f'prompt: {batch_prompt_embeds.dtype}')
         print(unet_added_conditions['time_ids'].dtype)
         print(unet_added_conditions['text_embeds'].dtype)