Update app.py

app.py

@@ -1,6 +1,8 @@
+
+
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
+from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler
 from huggingface_hub import hf_hub_download
 import spaces
 from PIL import Image
@@ -10,13 +12,8 @@ from translatepy import Translator
 translator = Translator()
 
 # Constants
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "tianweiy/DMD2"
-checkpoints = {
-    "1-Step" : ["dmd2_sdxl_1step_unet_fp16.bin", 1],
-    "4-Step" : ["dmd2_sdxl_4step_unet_fp16.bin", 4],
-}
-loaded = None
+model = "stabilityai/stable-diffusion-3-medium"
+vae_model = "madebyollin/sdxl-vae-fp16-fix"
 
 CSS = """
 .gradio-container {
@@ -35,37 +32,44 @@ JS = """function () {
 }"""
 
 
+# Load VAE component
+vae = AutoencoderKL.from_pretrained(
+    vae_model, 
+    torch_dtype=torch.float16
+)
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
-    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-    pipe = DiffusionPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
+    pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16).to("cuda")
+
+pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
 
 # Function 
 @spaces.GPU()
-def generate_image(prompt, ckpt="4-Step"):
-    global loaded
+def generate_image(
+    prompt,
+    negative="low quality",
+    width=1024,
+    height=1024,
+    scale=1.5,
+    steps=30,
+    clip=3):
     
     prompt = str(translator.translate(prompt, 'English'))
 
-    print(prompt)
+    print(f'prompt:{prompt}')
     
-    checkpoint = checkpoints[ckpt][0]
-    num_inference_steps = checkpoints[ckpt][1]
-
-    if loaded != num_inference_steps:
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-        pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, checkpoint), map_location="cuda"))
-        loaded = num_inference_steps
-
-    if loaded == 1:
-        timesteps=[399]
-    else:
-        timesteps=[999, 749, 499, 249]
-
-    results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=timesteps)
-    return results.images[0]
+    image = pipe(
+        prompt, 
+        negative_prompt=negative, 
+        width=width,
+        height=height,
+        guidance_scale=scale,
+        num_inference_steps=steps,
+        clip_skip=clip,
+    )
+    return image.images[0]
 
 
 examples = [
@@ -82,14 +86,53 @@ examples = [
 # Gradio Interface
 
 with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
-    gr.HTML("<h1><center>DMD2🦖</center></h1>")
-    gr.HTML("<p><center><a href='https://huggingface.co/tianweiy/DMD2'>DMD2</a> text-to-image generation</center><br><center>Multi-Languages, 4-step is higher quality & 2X slower</center></p>")
+    gr.HTML("<h1><center>Mobius💠</center></h1>")
+    gr.HTML("<p><center><a href='https://huggingface.co/Corcelio/mobius'>mobius</a> text-to-image generation</center><br><center>Multi-Languages. Adding default prompts to enhance.</center></p>")
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label='Enter Your Prompt', scale=8)
-            ckpt = gr.Dropdown(label='Steps',choices=['1-Step', '4-Step'], value='4-Step', interactive=True)
+            prompt = gr.Textbox(label='Enter Your Prompt', value="best quality, HD, aesthetic", scale=6)
             submit = gr.Button(scale=1, variant='primary')
-    img = gr.Image(label='DMD2 Generated Image')    
+    img = gr.Image(label='Mobius Generated Image')
+    with gr.Accordion("Advanced Options", open=False):
+        with gr.Row():
+            negative = gr.Textbox(label="Negative prompt", value="low quality")
+        with gr.Row():
+            width = gr.Slider(
+                label="Width",
+                minimum=512,
+                maximum=1280,
+                step=8,
+                value=1024,
+            )
+            height = gr.Slider(
+                label="Height",
+                minimum=512,
+                maximum=1280,
+                step=8,
+                value=1024,
+            )
+        with gr.Row():
+            scale = gr.Slider(
+                label="Guidance",
+                minimum=3.5,
+                maximum=7,
+                step=0.1,
+                value=7,
+            )
+            steps = gr.Slider(
+                label="Steps",
+                minimum=1,
+                maximum=50,
+                step=1,
+                value=50,
+            )
+            clip = gr.Slider(
+                label="Clip Skip",
+                minimum=1,
+                maximum=10,
+                step=1,
+                value=3,
+            )      
     gr.Examples(
         examples=examples,
         inputs=prompt,
@@ -99,11 +142,11 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
     )
 
     prompt.submit(fn=generate_image,
-                 inputs=[prompt, ckpt],
+                 inputs=[prompt, negative, width, height, scale, steps, clip],
                  outputs=img,
                  )
     submit.click(fn=generate_image,
-                 inputs=[prompt, ckpt],
+                 inputs=[prompt, negative, width, height, scale, steps, clip],
                  outputs=img,
                  )