Update app.py

app.py

@@ -5,25 +5,37 @@ from diffusers import StableDiffusionPipeline, LCMScheduler
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-adapter_id = "jasperai/flash-sd"
+
+transformer = Transformer2DModel.from_pretrained(
+  "PixArt-alpha/PixArt-XL-2-1024-MS",
+  subfolder="transformer",
+  torch_dtype=torch.float16
+)
+transformer = PeftModel.from_pretrained(
+  transformer,
+  "jasperai/flash-pixart"
+)
+
 
 if torch.cuda.is_available():
     torch.cuda.max_memory_allocated(device=device)
-    pipe = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-          use_safetensors=True,
+    pipe = PixArtAlphaPipeline.from_pretrained(
+      "PixArt-alpha/PixArt-XL-2-1024-MS",
+      transformer=transformer,
+      torch_dtype=torch.float16
     )
     pipe.enable_xformers_memory_efficient_attention()
     pipe = pipe.to(device)
 else: 
-    pipe = StableDiffusionPipeline.from_pretrained(
-      "runwayml/stable-diffusion-v1-5",
-      use_safetensors=True,
+    pipe = PixArtAlphaPipeline.from_pretrained(
+      "PixArt-alpha/PixArt-XL-2-1024-MS",
+      transformer=transformer,
+      torch_dtype=torch.float16
     )
     pipe = pipe.to(device)
 
 pipe.scheduler = LCMScheduler.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
+    "PixArt-alpha/PixArt-XL-2-1024-MS",
     subfolder="scheduler",
     timestep_spacing="trailing",
 )
@@ -32,7 +44,8 @@ pipe.load_lora_weights(adapter_id)
 pipe.fuse_lora()
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 512
+MAX_IMAGE_SIZE = 1024
+NUM_INFERENCE_STEPS = 4
 
 def infer(prompt, seed, randomize_seed, num_inference_steps):
 
@@ -59,7 +72,7 @@ examples = [
 css="""
 #col-container {
     margin: 0 auto;
-    max-width: 520px;
+    max-width: 512px;
 }
 """
 
@@ -104,17 +117,6 @@ with gr.Blocks(css=css) as demo:
             
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             
-
-            
-            with gr.Row():
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=2,
-                    maximum=8,
-                    step=1,
-                    value=4,
-                )
         
         gr.Examples(
             examples = examples,
@@ -123,7 +125,7 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, seed, randomize_seed, num_inference_steps],
+        inputs = [prompt, seed, randomize_seed, NUM_INFERENCE_STEPS],
         outputs = [result]
     )