Update app.py

app.py

@@ -1,76 +1,100 @@
 from __future__ import annotations
-
 import os
 import random
 import uuid
-from typing import Tuple
-
 import gradio as gr
+import spaces
 import numpy as np
+import uuid
+from diffusers import PixArtAlphaPipeline, LCMScheduler
 import torch
-from diffusers import LCMScheduler, PixArtAlphaPipeline
+from typing import Tuple
+from datetime import datetime
 
-# Use a more descriptive variable name
-MODEL_NAME = "PixArt-alpha/PixArt-LCM-XL-2-1024-MS"
 
-# Move environment variable checks and definitions to the top for better readability
-DESCRIPTION = """# Instant Image
+DESCRIPTION = """ # Instant Image
         ### Super fast text to Image Generator.
         ### <span style='color: red;'>You may change the steps from 4 to 8, if you didn't get satisfied results.
         ### First Image processing takes time then images generate faster.
         """
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
+MAX_SEED = np.iinfo(np.int32).max
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4192"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 PORT = int(os.getenv("DEMO_PORT", "15432"))
 
-# Check CUDA availability early on
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
-
-# Cache examples only if CUDA is available
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
-
-MAX_SEED = np.iinfo(np.int32).max
-NUM_IMAGES_PER_PROMPT = 1
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Use Enum for better style management
-from enum import Enum
 
-class Style(Enum):
-    NO_STYLE = ("(No style)", "{prompt}", "")
-    CINEMATIC = ("Cinematic", "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy", "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured")
-    REALISTIC = ("Realistic", "Photorealistic {prompt} . Ulta-realistic, professional, 4k, highly detailed", "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, disfigured")
-    ANIME = ("Anime", "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed", "photo, deformed, black and white, realism, disfigured, low contrast")
-    DIGITAL_ART = ("Digital Art", "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed", "photo, photorealistic, realism, ugly")
-    PIXEL_ART = ("Pixel art", "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics", "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic")
-    FANTASY_ART = ("Fantasy art", "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy", "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white")
-    THREE_D_MODEL = ("3D Model", "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting", "ugly, deformed, noisy, low poly, blurry, painting")
+style_list = [
+    {
+        "name": "(No style)",
+        "prompt": "{prompt}",
+        "negative_prompt": "",
+    },
+    {
+        "name": "Cinematic",
+        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
+    },
+    {
+        "name": "Realistic",
+        "prompt": "Photorealistic {prompt} . Ulta-realistic, professional, 4k, highly detailed",
+        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, disfigured",
+    },
+    {
+        "name": "Anime",
+        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
+        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
+    },
+    {
+        "name": "Digital Art",
+        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
+        "negative_prompt": "photo, photorealistic, realism, ugly",
+    },
+    {
+        "name": "Pixel art",
+        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
+        "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
+    },
+    {
+        "name": "Fantasy art",
+        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
+    },
+]
 
-    def __init__(self, name, prompt, negative_prompt):
-        self.name = name
-        self.prompt = prompt
-        self.negative_prompt = negative_prompt
 
-# Use the Enum values directly
-styles = {style.name: (style.prompt, style.negative_prompt) for style in Style}
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
 STYLE_NAMES = list(styles.keys())
-DEFAULT_STYLE_NAME = Style.NO_STYLE.name
+DEFAULT_STYLE_NAME = "(No style)"
+NUM_IMAGES_PER_PROMPT = 1
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+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
 
-# Load the pipeline only if CUDA is available
 if torch.cuda.is_available():
+
     pipe = PixArtAlphaPipeline.from_pretrained(
-        MODEL_NAME,
+        "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
         torch_dtype=torch.float16,
         use_safetensors=True,
     )
 
     if os.getenv('CONSISTENCY_DECODER', False):
         print("Using DALL-E 3 Consistency Decoder")
-        # Assuming ConsistencyDecoderVAE is defined elsewhere
         pipe.vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
 
     if ENABLE_CPU_OFFLOAD:
@@ -78,8 +102,8 @@ if torch.cuda.is_available():
     else:
         pipe.to(device)
         print("Loaded on Device!")
-
-    # Speed-up T5
+        
+    # speed-up T5
     pipe.text_encoder.to_bettertransformer()
 
     if USE_TORCH_COMPILE:
@@ -87,40 +111,37 @@ if torch.cuda.is_available():
         print("Model Compiled!")
 
 def save_image(img):
-    # Generate image names in a temporary directory
-    os.makedirs("tmp", exist_ok=True)
-    unique_name = os.path.join("tmp", f"{uuid.uuid4()}.png")
+    unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
 
+
+
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    return random.randint(0, MAX_SEED) if randomize_seed else seed
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
 
-# No need to use @spaces.GPU if you're checking CUDA availability within the function
+@spaces.GPU(duration=30)
 def generate(
-    prompt: str,
-    negative_prompt: str = "",
-    style: str = DEFAULT_STYLE_NAME,
-    use_negative_prompt: bool = False,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    inference_steps: int = 8,
-    randomize_seed: bool = False,
-    use_resolution_binning: bool = True,
-    progress=gr.Progress(track_tqdm=True),
+        prompt: str,
+        negative_prompt: str = "",
+        style: str = DEFAULT_STYLE_NAME,
+        use_negative_prompt: bool = False,
+        seed: int = 0,
+        width: int = 1024,
+        height: int = 1024,
+        inference_steps: int = 8,
+        randomize_seed: bool = False,
+        use_resolution_binning: bool = True,
+        progress=gr.Progress(track_tqdm=True),
 ):
-    if not torch.cuda.is_available():
-        return "This demo requires a GPU to run.", seed
-
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
-
+    
     if not use_negative_prompt:
-        negative_prompt = None 
-
-    prompt, negative_prompt = styles.get(style, styles[DEFAULT_STYLE_NAME])
-    prompt = prompt.replace("{prompt}", prompt)
+        negative_prompt = None  # type: ignore
+    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
 
     images = pipe(
         prompt=prompt,
@@ -136,6 +157,7 @@ def generate(
     ).images
 
     image_paths = [save_image(img) for img in images]
+    print(image_paths)
     return image_paths, seed
 
 
@@ -170,17 +192,17 @@ with gr.Blocks(css=css) as demo:
                 )
                 run_button = gr.Button("Run", scale=0)
             result = gr.Gallery(label="Result", columns=1,  show_label=False)
-
+                    
     with gr.Accordion("Advanced options", open=False):
         with gr.Group():
             with gr.Row():
                 use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False, visible=True)
                 negative_prompt = gr.Text(
-                    label="Negative prompt",
-                    max_lines=1,
-                    placeholder="Enter a negative prompt",
-                    visible=True,
-                )
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=True,
+            )
             style_selection = gr.Radio(
                 show_label=True,
                 container=True,
@@ -195,7 +217,7 @@ with gr.Blocks(css=css) as demo:
                 maximum=MAX_SEED,
                 step=1,
                 value=0,
-            )
+                )
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             with gr.Row(visible=True):
                 width = gr.Slider(
@@ -242,7 +264,7 @@ with gr.Blocks(css=css) as demo:
             run_button.click,
         ],
         batch=True,
-        max_batch_size=10,
+        max_batch_size=10, 
         fn=generate,
         inputs=[
             prompt,
@@ -260,4 +282,4 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=200).launch(server_port=PORT)
+    demo.queue(max_size=200).launch()