Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
+from diffusers import StableDiffusionXLPipeline, AutoencoderKL
 from huggingface_hub import hf_hub_download
 import spaces
 from PIL import Image
@@ -10,13 +10,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
+base = "Corcelio/openvision"
+vae = "madebyollin/sdxl-vae-fp16-fix"
 
 CSS = """
 .gradio-container {
@@ -34,38 +29,42 @@ JS = """function () {
   }
 }"""
 
-
+# Load VAE component
+vae = AutoencoderKL.from_pretrained(
+    "madebyollin/sdxl-vae-fp16-fix", 
+    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(base, vae=vae, torch_dtype=torch.float16).to("cuda")
+
 
 
 # Function 
 @spaces.GPU()
-def generate_image(prompt, ckpt="4-Step"):
-    global loaded
+def generate_image(
+    prompt,
+    negative,
+    width=1024,
+    height=1024,
+    scale=1.5,
+    steps=30):
     
     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,
+    )
+    print(image.images[0])
+    return image.images[0]
 
 
 examples = [
@@ -82,14 +81,46 @@ examples = [
 # Gradio Interface
 
 with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
-    gr.HTML("<h1><center>Adobe 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>OpenVision</center></h1>")
+    gr.HTML("<p><center><a href='https://huggingface.co/Corcelio/openvision'>OpenVision</a> text-to-image generation</center><br><center>Multi-Languages. Midjourney Aesthetic for All Your Images</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', scale=6)
             submit = gr.Button(scale=1, variant='primary')
-    img = gr.Image(label='DMD2 Generated Image')    
+    img = gr.Image(label='OpenVision 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 Scale",
+                minimum=0,
+                maximum=50,
+                step=0.1,
+                value=1.5,
+            )
+            steps = gr.Slider(
+                label="Steps",
+                minimum=1,
+                maximum=50,
+                step=1,
+                value=30,
+            )            
     gr.Examples(
         examples=examples,
         inputs=prompt,
@@ -99,11 +130,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],
                  outputs=img,
                  )
     submit.click(fn=generate_image,
-                 inputs=[prompt, ckpt],
+                 inputs=[prompt, negative, width, height, scale, steps],
                  outputs=img,
                  )