Update app.py

app.py

@@ -5,10 +5,11 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline, AutoencoderKL, UNet2DConditionModel, FlowMatchEulerDiscreteScheduler
 from transformers import AutoTokenizer, AutoModel
+from diffusers import DiffusionPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-model_repo_id = "AiArtLab/sdxs3d"
+model_repo_id = "AiArtLab/sdxs"
 max_length = 150
 
 class SimpleDiffusionPipeline(DiffusionPipeline):
@@ -94,18 +95,25 @@ class SimpleDiffusionPipeline(DiffusionPipeline):
 
         return images
 
-vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae", torch_dtype=dtype).to(device)
-unet = UNet2DConditionModel.from_pretrained(model_repo_id, subfolder="unet", torch_dtype=dtype).to(device)
-tokenizer = AutoTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
-text_encoder = AutoModel.from_pretrained(model_repo_id, subfolder="text_encoder", torch_dtype=dtype).to(device)
-scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler")
-
-pipe = SimpleDiffusionPipeline(
-    vae=vae,
-    text_encoder=text_encoder,
-    tokenizer=tokenizer,
-    unet=unet,
-    scheduler=scheduler,
+#vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae", torch_dtype=dtype).to(device)
+#unet = UNet2DConditionModel.from_pretrained(model_repo_id, subfolder="unet", torch_dtype=dtype).to(device)
+#tokenizer = AutoTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
+#text_encoder = AutoModel.from_pretrained(model_repo_id, subfolder="text_encoder", torch_dtype=dtype).to(device)
+#scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler")
+
+#pipe = SimpleDiffusionPipeline(
+#    vae=vae,
+#    text_encoder=text_encoder,
+#    tokenizer=tokenizer,
+#    unet=unet,
+#    scheduler=scheduler,
+#).to(device)
+
+pipe_id = "AiArtLab/sdxs"
+pipe = SdxsPipeline.from_pretrained(
+    pipe_id,
+    torch_dtype=dtype,
+    trust_remote_code=True
 ).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max