Update app.py

app.py

@@ -11,6 +11,8 @@ from transformers import pipeline, set_seed
 from diffusers import DiffusionPipeline, StableDiffusionPipeline
 
 
+# Config
+DEVICE = "cuda"
 
 # GPT2
 generator = pipeline('text-generation', model='gpt2')
@@ -19,15 +21,19 @@ set_seed(42)
 generator("Hello world, I'm vizard,", max_length=50, num_return_sequences=3)
 
 # SD v1.4
-pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True, revision="fp16", torch_dtype=torch.float16)
-pipe = pipe.to(device)
-torch.backends.cudnn.benchmark = True
+def get_stable_diffusion_v14_pipeline():
+  model_id = "CompVis/stable-diffusion-v1-4"
+  pipeline = StableDiffusionPipeline.from_pretrained(mode_id)
+  # pipeline = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True, revision="fp16", torch_dtype=torch.float16)
+  pipeline = pipeline.to(DEVICE)
+  torch.backends.cudnn.benchmark = True
+  return pipeline
 
 # SD v1.5
 def get_stable_diffusion_v15_pipeline():
   model_id = "runwayml/stable-diffusion-v1-5"
   pipeline = DiffusionPipeline.from_pretrained(mode_id)
-  pipeline = pipeline.to(device)
+  pipeline = pipeline.to(DEVICE)
   return pipeline
   
 # main