Update app.py

app.py

@@ -6,10 +6,11 @@ from PIL import Image, ImageOps
 
 
 # Load pipeline once
-model_id = 'Deci/DeciDiffusion-v1-0'
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, custom_pipeline=model_id, torch_dtype=torch.float32)
-pipe.unet = pipe.unet.from_pretrained(model_id, subfolder='flexible_unet', torch_dtype=torch.float32)
+checkpoint = "Deci/DeciDiffusion-v2-0"
+
+pipe = StableDiffusionPipeline.from_pretrained(checkpoint, custom_pipeline=checkpoint, torch_dtype=torch.float32)
+pipe.unet = pipe.unet.from_pretrained(checkpoint, subfolder='flexible_unet', torch_dtype=torch.float32)
 pipe = pipe.to(device)
 
 
@@ -22,13 +23,14 @@ def read_content(file_path: str) -> str:
     return content
 
 
-def predict(_prompt: str, _steps: int = 30, _seed: int = 42, _guidance_scale: float = 7.5, _negative_prompt: str = ""):
+def predict(_prompt: str, _seed: int = 42, _guidance_scale: float = 7.5, _guidance_rescale: float = 0.7, _negative_prompt: str = ""):
     _negative_prompt = [_negative_prompt] if _negative_prompt else None
 
     output = pipe(prompt=[_prompt],
                   negative_prompt=_negative_prompt,
-                  num_inference_steps=int(_steps),
+                  num_inference_steps=16,
                   guidance_scale=_guidance_scale,
+                  guidance_rescale=_guidance_rescale,
                   generator=torch.Generator(device).manual_seed(_seed),
                   )
     output_image = output.images[0]
@@ -79,9 +81,10 @@ with demo:
 
             with gr.Accordion(label="Advanced Settings", open=False):
                 with gr.Row(mobile_collapse=False, equal_height=True):
-                    steps = gr.Slider(value=30, minimum=15, maximum=50, step=1, label="steps", interactive=True)
                     seed = gr.Slider(value=42, minimum=1, maximum=100, step=1, label="seed", interactive=True)
-                    guidance_scale = gr.Slider(value=7.5, minimum=1, maximum=15, step=0.1, label='guidance_scale', interactive=True)
+                    guidance_scale = gr.Slider(value=7.5, minimum=2, maximum=20, step=0.1, label='guidance_scale', interactive=True)
+                    guidance_rescale = gr.Slider(value=0.7, minimum=0.0, maximum=0.99, step=0.05, label='guidance_rescale', interactive=True)
+
 
                 with gr.Row(mobile_collapse=False, equal_height=True):
                     negative_prompt = gr.Textbox(label="negative_prompt", placeholder="Your negative prompt",
@@ -93,7 +96,7 @@ with demo:
             image_out = gr.Image(label="Output", elem_id="output-img", height=400)
 
     btn.click(fn=predict,
-              inputs=[prompt, steps, seed, guidance_scale, negative_prompt],
+              inputs=[prompt, seed, guidance_scale, guidance_rescale, negative_prompt],
               outputs=[image_out],
               api_name='run')