Update app.py

app.py

@@ -1,11 +1,17 @@
 import torch
-from diffusers import PixArtSigmaPipeline
+from pipeline import PixArtSigmaPipeline
 import gradio as gr
 import spaces
 
 # Load the pre-trained diffusion model
-pipe = PixArtSigmaPipeline.from_pretrained('ptx0/pixart-900m-1024-ft', torch_dtype=torch.bfloat16)
-pipe.to('cuda')
+base_model = "ptx0/pixart-900m-1024-ft-v0.7-stage1"
+stg2_model = "ptx0/pixart-900m-1024-ft-v0.7-stage2"
+torch_device = "cuda"
+base_pipeline = PixArtSigmaPipeline.from_pretrained(
+	base_model, use_safetensors=True
+).to(dtype=torch_precision, device=torch_device)
+stg2_pipeline = PixArtSigmaPipeline.from_pretrained(stg2_model, **base_pipeline.components)
+stg2_pipeline.transformer = PixArtTransformer2DModel.from_pretrained(stg2_model, subfolder="transformer").to(dtype=torch_precision, device=torch_device)
 import re
 
 def extract_resolution(resolution_str):
@@ -21,12 +27,29 @@ def extract_resolution(resolution_str):
 @spaces.GPU
 def generate(prompt, guidance_scale, num_inference_steps, resolution, negative_prompt):
     width, height = extract_resolution(resolution) or (1024, 1024)
-    return pipe(
-        prompt,
+    mixture_generator = torch.Generator().manual_seed(444)
+    stage1_strength = 0.6
+    latent_images = base_pipeline(
+        prompt=prompt,
         negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=1,
+        generator=mixture_generator,
         guidance_scale=guidance_scale,
+        output_type="latent",
+        denoising_end=stage1_strength,
+        width=width,
+        height=height
+    ).images
+    return refined_images = stg2_pipeline(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        latents=latent_images,
         num_inference_steps=num_inference_steps,
-        width=width, height=height
+        num_images_per_prompt=1,
+        generator=mixture_generator,
+        guidance_scale=guidance_scale,
+        denoising_start=stage1_strength
     ).images
 
 # Example prompts to demonstrate the model's capabilities