Update app.py

app.py

@@ -1,10 +1,10 @@
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusion3Pipeline
 import gradio as gr
 import spaces
 
 # Load the pre-trained diffusion model
-pipe = DiffusionPipeline.from_pretrained('ptx0/pseudo-flex-v2', torch_dtype=torch.bfloat16)
+pipe = StableDiffusion3Pipeline.from_pretrained('ptx0/sd3-diffusion-vpred-zsnr', torch_dtype=torch.bfloat16)
 pipe.to('cuda')
 import re
 
@@ -25,7 +25,7 @@ def generate(prompt, guidance_scale, guidance_rescale, num_inference_steps, reso
         prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
-        guidance_rescale=guidance_rescale,
+        #guidance_rescale=guidance_rescale,
         num_inference_steps=num_inference_steps,
         width=width, height=height
     ).images
@@ -50,6 +50,6 @@ iface = gr.Interface(
     ],
     outputs=gr.Gallery(height=1024, min_width=1024, columns=2),
     examples=example_prompts,
-    title="Flex v2 (SD 2.1-v) Demonstration",
-    description="Flex v2 is a multi-aspect finetune of SD 2.1-v (768px) that is up-sized to a base resolution of 1 megapixel (1024px). This model utilises a zero-terminal SNR noise schedule, formulated to allow for very dark and very bright images."
+    title="SD3 Diffusion Demonstration",
+    description="Stable Diffusion 3 Diffusion is a v-prediction model trained to eliminate the rectified flow schedule from Stable Diffusion 3 as an experiment into this model architecture and its parameterisations."
 ).launch()