Update app.py

app.py

@@ -1,19 +1,14 @@
-#!/usr/bin/env python
 from __future__ import annotations
-import argparse
 import os
-import sys
 import random
+import uuid
 import gradio as gr
 import numpy as np
 import uuid
-import spaces
-from diffusers import ConsistencyDecoderVAE, DPMSolverMultistepScheduler, Transformer2DModel, AutoencoderKL
+from diffusers import PixArtAlphaPipeline, LCMScheduler
 import torch
 from typing import Tuple
 from datetime import datetime
-from peft import PeftModel
-from diffusers_patches import pixart_sigma_init_patched_inputs, PixArtSigmaPipeline
 
 
 DESCRIPTION = """ # Instant Image
@@ -91,36 +86,13 @@ style_list = [
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
 STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = "(No style)"
-SCHEDULE_NAME = ["DPM-Solver"]
-DEFAULT_SCHEDULE_NAME = "DPM-Solver"
 NUM_IMAGES_PER_PROMPT = 1
 
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    if not negative:
-        negative = ""
-    return p.replace("{prompt}", positive), n + negative
-
-
 if torch.cuda.is_available():
-    weight_dtype = torch.float16
-    T5_token_max_length = 300
 
-    # tmp patches for diffusers PixArtSigmaPipeline Implementation
-    print(
-        "Changing _init_patched_inputs method of diffusers.models.Transformer2DModel "
-        "using scripts.diffusers_patches.pixart_sigma_init_patched_inputs")
-    setattr(Transformer2DModel, '_init_patched_inputs', pixart_sigma_init_patched_inputs)
-
-    transformer = Transformer2DModel.from_pretrained(
-        "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
-        subfolder='transformer',
-        torch_dtype=weight_dtype,
-    )
-    pipe = PixArtSigmaPipeline.from_pretrained(
-        "PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
-        transformer=transformer,
-        torch_dtype=weight_dtype,
+    pipe = PixArtAlphaPipeline.from_pretrained(
+        "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
+        torch_dtype=torch.float16,
         use_safetensors=True,
     )
 
@@ -141,7 +113,6 @@ if torch.cuda.is_available():
         pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True)
         print("Model Compiled!")
 
-
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
@@ -152,10 +123,6 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
-@torch.no_grad()
-@torch.inference_mode()
-@spaces.GPU(duration=30)
 def generate(
         prompt: str,
         negative_prompt: str = "",
@@ -163,11 +130,9 @@ def generate(
         use_negative_prompt: bool = False,
         num_imgs: int = 1,
         seed: int = 0,
-        width: int = 2560,
-        height: int = 2560,
-        schedule: str = 'DPM-Solver',
-        dpms_guidance_scale: float = 3.5,
-        dpms_inference_steps: int = 9,
+        width: int = 1024,
+        height: int = 1024,
+        num_inference_steps: int = 4,
         randomize_seed: bool = False,
         use_resolution_binning: bool = True,
         progress=gr.Progress(track_tqdm=True),
@@ -175,15 +140,7 @@ def generate(
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
 
-    if schedule == 'DPM-Solver':
-        if not isinstance(pipe.scheduler, DPMSolverMultistepScheduler):
-            pipe.scheduler = DPMSolverMultistepScheduler()
-        num_inference_steps = dpms_inference_steps
-        guidance_scale = dpms_guidance_scale
-    else:
-        raise ValueError(f"Unknown schedule: {schedule}")
-
-    if not use_negative_prompt:
+if not use_negative_prompt:
         negative_prompt = None  # type: ignore
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
 
@@ -195,10 +152,9 @@ def generate(
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         generator=generator,
-        num_images_per_prompt=num_imgs,
+        num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
         use_resolution_binning=use_resolution_binning,
         output_type="pil",
-        max_sequence_length=T5_token_max_length,
     ).images
 
     image_paths = [save_image(img) for img in images]
@@ -242,16 +198,13 @@ with gr.Blocks(css="style.css") as demo:
         with gr.Group():
             with gr.Row():
                 use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False, visible=True)
-            with gr.Row(visible=True):
-                schedule = gr.Radio(
-                    show_label=True,
-                    container=True,
-                    interactive=True,
-                    choices=SCHEDULE_NAME,
-                    value=DEFAULT_SCHEDULE_NAME,
-                    label="Sampler Schedule",
-                    visible=True,
-                )
+                 negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=True,
+            )
+              
                 num_imgs = gr.Slider(
                     label="Num Images",
                     minimum=1,
@@ -287,29 +240,23 @@ with gr.Blocks(css="style.css") as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=2560,
+                    value=1024,
                 )
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=2560,
-                )
-            with gr.Row():
-                dpms_guidance_scale = gr.Slider(
-                    label="Temprature",
-                    minimum=3,
-                    maximum=4,
-                    step=0.1,
-                    value=3.5,
+                    value=1024,
                 )
-                dpms_inference_steps = gr.Slider(
-                    label="Steps",
-                    minimum=5,
-                    maximum=25,
-                    step=1,
-                    value=9,
+        with gr.Row():
+            inference_steps = gr.Slider(
+                label="Steps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=6,
+            )  value=9,
                 )
 
     gr.Examples(