Update app.py

app.py

@@ -4,6 +4,7 @@ import random
 from diffusers import AuraFlowPipeline
 import torch
 import spaces
+import PixArtSigmaPipeline
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -14,9 +15,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 #torch._inductor.config.epilogue_fusion = False
 #torch._inductor.config.coordinate_descent_check_all_directions = True
 
-pipe = AuraFlowPipeline.from_pretrained(
-	"fal/AuraFlow",
-    torch_dtype=torch.float16
+pipe = PixArtSigmaPipeline.from_pretrained(
+    "dataautogpt3/PixArt-Sigma-900M", 
+    torch_dtype=torch.float16,
 ).to("cuda")
 
 #pipe.transformer.to(memory_format=torch.channels_last)
@@ -62,17 +63,12 @@ css="""
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # AuraFlow 0.1
-        Demo of the [AuraFlow 0.1](https://huggingface.co/fal/AuraFlow) 6.8B parameters open source diffusion transformer model
+        # PixArt Sigma 900M
+        Demo of the [PixArt Sigma 900M](https://huggingface.co/fal/AuraFlow) 6.8B parameters open source diffusion transformer model
         [[blog](https://blog.fal.ai/auraflow/)] [[model](https://huggingface.co/fal/AuraFlow)] [[fal](https://fal.ai/models/fal-ai/aura-flow)]
         """)