Update app.py

app.py

@@ -46,10 +46,11 @@ class ModelWrapper:
         self.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
         self.alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
         self.num_step = num_step
-
+        
+    @spaces.GPU(enable_queue=True)
     def create_generator(self, model_id, checkpoint_path):
         generator = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(self.DTYPE)
-        state_dict = torch.load(checkpoint_path, map_location="cpu")
+        state_dict = torch.load(checkpoint_path, map_location="cuda")
         generator.load_state_dict(state_dict, strict=True)
         generator.requires_grad_(False)
         return generator
@@ -108,7 +109,7 @@ class ModelWrapper:
         eval_images = ((eval_images + 1.0) * 127.5).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1)
         return eval_images 
         
-    @spaces.GPU(enable_queue=True)
+    
     @torch.no_grad()
     def inference(self, prompt, seed, height, width, num_images, fast_vae_decode):
         print("Running model inference...")
@@ -196,9 +197,6 @@ def create_demo():
     num_step = 4
     revision = None
 
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True 
-
     accelerator = Accelerator()
 
     model = ModelWrapper(model_id, checkpoint_path, precision, image_resolution, latent_resolution, num_train_timesteps, conditioning_timestep, num_step, revision, accelerator)
@@ -211,10 +209,10 @@ def create_demo():
                 run_button = gr.Button("Run")
                 with gr.Accordion(label="Advanced options", open=True):
                     seed = gr.Slider(label="Seed", minimum=-1, maximum=1000000, step=1, value=0)
-                    num_images = gr.Slider(label="Number of generated images", minimum=1, maximum=16, step=1, value=16)
+                    num_images = gr.Slider(label="Number of generated images", minimum=1, maximum=16, step=1, value=1)
                     fast_vae_decode = gr.Checkbox(label="Use Tiny VAE for faster decoding", value=True)
-                    height = gr.Slider(label="Image Height", minimum=512, maximum=1536, step=64, value=1024)
-                    width = gr.Slider(label="Image Width", minimum=512, maximum=1536, step=64, value=1024)
+                    height = gr.Slider(label="Image Height", minimum=512, maximum=1536, step=64, value=512)
+                    width = gr.Slider(label="Image Width", minimum=512, maximum=1536, step=64, value=512)
             with gr.Column():
                 result = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=1024)
                 error_message = gr.Text(label="Job Status")