Update app.py

app.py

@@ -4,18 +4,30 @@ import gradio as gr
 import spaces
 
 # Load the pre-trained diffusion model
-pipe = DiffusionPipeline.from_pretrained('ptx0/terminus-xl-velocity-v2', torch_dtype=torch.bfloat16)
+pipe = DiffusionPipeline.from_pretrained('ptx0/pseudo-flex-v2', torch_dtype=torch.bfloat16)
 pipe.to('cuda')
+import re
+
+def extract_resolution(resolution_str):
+    match = re.match(r'(\d+)x(\d+)', resolution_str)
+    if match:
+        width = int(match.group(1))
+        height = int(match.group(2))
+        return (width, height)
+    else:
+        return None
 
 # Define the image generation function with adjustable parameters and a progress bar
 @spaces.GPU
-def generate(prompt, guidance_scale, guidance_rescale, num_inference_steps, negative_prompt):
+def generate(prompt, guidance_scale, guidance_rescale, num_inference_steps, resolution, negative_prompt):
+    width, height = extract_resolution(resolution) or (1024, 1024)
     return pipe(
         prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         guidance_rescale=guidance_rescale,
-        num_inference_steps=num_inference_steps
+        num_inference_steps=num_inference_steps,
+        width=width, height=height
     ).images
 
 # Example prompts to demonstrate the model's capabilities
@@ -25,18 +37,19 @@ example_prompts = [
     ["An abstract painting of joy and energy in bright colors", 9.0, 30, "dark, dull"]
 ]
 
-# Create a Gradio interface
+# Create a Gradio interface, 1024x1024,1152x960,896x1152
 iface = gr.Interface(
     fn=generate,
     inputs=[
         gr.Text(label="Enter your prompt"),
-        gr.Slider(1, 20, step=0.1, label="Guidance Scale", value=11.5),
+        gr.Slider(1, 20, step=0.1, label="Guidance Scale", value=11.5)
         gr.Slider(0, 1, step=0.1, label="Rescale classifier-free guidance", value=0.7),
         gr.Slider(1, 50, step=1, label="Number of Inference Steps", value=25),
+        gr.Radio(["1024x1024", "1152x960", "896x1152"], label="Resolution", value="1152x960"),
         gr.Text(value="underexposed, blurry, ugly, washed-out", label="Negative Prompt")
     ],
     outputs=gr.Gallery(height=1024, min_width=1024, columns=2),
     examples=example_prompts,
-    title="Terminus XL Velocity v2.0 Demonstration",
-    description="Terminus XL is a v-prediction model trained with a zero-terminal SNR noise schedule, allowing it to create very dark or very bright images. Occasionally, it will work pretty well for typography, eg. text in images."
+    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."
 ).launch()