Send model to GPU

app.py

@@ -5,6 +5,7 @@ from diffusers import StableDiffusionPipeline
 
 model_id = "hakurei/waifu-diffusion"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision='fp16')
+pipe = pipe.to("cuda")
 num_samples = 2
 
 def infer(prompt):
@@ -128,6 +129,9 @@ css = """
         #prompt-container{
             gap: 0;
         }
+        #generated_id{
+            min-height: 700px
+        }
 """
 block = gr.Blocks(css=css)
 with block as demo:
@@ -169,7 +173,7 @@ with block as demo:
                     rounded=(False, True, True, False),
                 )
                
-        gallery = gr.Gallery(label="Generated images", show_label=False).style(
+        gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="generated_id").style(
             grid=[2], height="auto"
         )
         text.submit(infer, inputs=[text], outputs=gallery)