multimodalart HF staff commited on
Commit
70833cb
·
1 Parent(s): 78afb32

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -4,11 +4,13 @@ import PIL.Image
4
  import numpy as np
5
  import random
6
  import torch
 
7
 
8
  ldm_pipeline = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
9
 
10
  def predict(prompt, steps=100, seed=42, guidance_scale=6.0):
11
  torch.cuda.empty_cache()
 
12
  generator = torch.manual_seed(seed)
13
  images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["sample"]
14
  return images[0]
 
4
  import numpy as np
5
  import random
6
  import torch
7
+ import subprocess
8
 
9
  ldm_pipeline = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
10
 
11
  def predict(prompt, steps=100, seed=42, guidance_scale=6.0):
12
  torch.cuda.empty_cache()
13
+ print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
14
  generator = torch.manual_seed(seed)
15
  images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["sample"]
16
  return images[0]