Update app.py

app.py

@@ -94,13 +94,23 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(device)
 
 
-model_dict = {
-    'Blazing Drive V13md': 'ironjr/BlazingDriveV13md',
-    # 'Real Cartoon Pixar V5': 'ironjr/RealCartoon-PixarV5',
-}
+if opt.model is None:
+    model_dict = {
+        # 'Blazing Drive V11m': 'ironjr/BlazingDriveV11m',
+        # 'Real Cartoon Pixar V5': 'ironjr/RealCartoon-PixarV5',
+        'Kohaku V2.1': 'KBlueLeaf/kohaku-v2.1',
+        # 'Realistic Vision V5.1': 'ironjr/RealisticVisionV5-1',
+        # 'Stable Diffusion V1.5': 'runwayml/stable-diffusion-v1-5',
+    }
+else:
+    if opt.model.endswith('.safetensors'):
+        opt.model = os.path.abspath(os.path.join('checkpoints', opt.model))
+    model_dict = {os.path.splitext(os.path.basename(opt.model))[0]: opt.model}
+
+dtype = torch.float32 if device == 'cpu' else torch.float16
 
 models = {
-    k: SemanticDrawPipeline(device, sd_version='1.5', hf_key=v).cuda()
+    k: SemanticDrawPipeline(device, dtype=dtype, sd_version='1.5', hf_key=v, has_i2t=False)
     for k, v in model_dict.items()
 }