update app.py

app.py

@@ -3,21 +3,29 @@ import numpy as np
 import random
 
 # import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline, FluxPipeline
+from diffusers import DiffusionPipeline #, FluxPipeline
 import torch
-import os
+# import os
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-model_repo = "black-forest-labs/FLUX.1-dev"  # Replace to the model you would like to use
-
-pipe = FluxPipeline.from_pretrained(model_repo, 
-                                    torch_dtype=torch_dtype, 
-                                    token=f"{os.environ.get('tkn')}",
-                                    #force_download=True,
-                                    ).to(device)
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+# model_repo = "black-forest-labs/FLUX.1-dev"  # Replace to the model you would like to use
+
+# pipe = FluxPipeline.from_pretrained(model_repo, 
+#                                     torch_dtype=torch_dtype, 
+#                                     token=f"{os.environ.get('tkn')}",
+#                                     #force_download=True,
+#                                     ).to(device)
 #pipe.enable_model_cpu_offload()
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+model_repo = "stabilityai/sdxl-turbo"
+pipe = DiffusionPipeline.from_pretrained(
+    model_repo,
+    torch_dtype=torch_dtype
+).to(device)
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024