Update app.py

app.py

@@ -3,7 +3,6 @@
 import os
 import random
 import uuid
-import json
 
 import gradio as gr
 import numpy as np
@@ -13,20 +12,32 @@ import torch
 from diffusers import StableDiffusion3Pipeline, DPMSolverMultistepScheduler, AutoencoderKL
 from huggingface_hub import snapshot_download
 
-# BaseConditions
-bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
-bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
-default_negative = os.getenv("default_negative","")
+huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+
+model_path = snapshot_download(
+    repo_id="stabilityai/stable-diffusion-3-medium", 
+    revision="refs/pr/26",
+    repo_type="model", 
+    ignore_patterns=["*.md", "*.gitattributes"],
+    local_dir="stable-diffusion-3-medium",
+    token=huggingface_token, # yeni bir token-id yazın.
+)
+
+DESCRIPTION = """# Stable Diffusion 3"""
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
 
-def check_text(prompt, negative=""):
-    for i in bad_words:
-        if i in prompt:
-            return True
-    for i in bad_words_negative:
-        if i in negative:
-            return True
-    return False
+MAX_SEED = np.iinfo(np.int32).max
+CACHE_EXAMPLES = False
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
+USE_TORCH_COMPILE = False
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+pipe = StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=torch.float16)
 
+# Define styles and collage templates
 style_list = [
     {
         "name": "3840 x 2160",
@@ -127,48 +138,10 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
         negative = ""
     return p.replace("{prompt}", positive), n + negative
 
-DESCRIPTION = """"""
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
-
-MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-
-# Download the model to a local directory
-model_path = snapshot_download(
-    repo_id="stabilityai/stable-diffusion-3-medium", 
-    revision="refs/pr/26",
-    repo_type="model", 
-    ignore_patterns=["*.md", "*.gitattributes"],
-    local_dir="stable-diffusion-3-medium",
-    token=huggingface_token,
-)
-
-if torch.cuda.is_available():
-    pipe = StableDiffusion3Pipeline.from_pretrained(
-        model_path,
-        torch_dtype=torch.float16,
-    ).to(device)
-
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-    else:
-        pipe.to(device)
-        print("Loaded on Device!")
-
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
-
-def save_image(img, path):
-    img.save(path)
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -182,82 +155,77 @@ def generate(
     use_negative_prompt: bool = False,
     style: str = DEFAULT_STYLE_NAME,
     collage_style: str = DEFAULT_COLLAGE_STYLE_NAME,
-    grid_size: str = "2x2",
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
-    guidance_scale: float = 3,
+    guidance_scale: float = 7,
     randomize_seed: bool = False,
+    num_inference_steps=30,
+    NUM_IMAGES_PER_PROMPT=1,
     use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if check_text(prompt, negative_prompt):
-        raise ValueError("Prompt contains restricted words.")
+    pipe.to(device)
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator().manual_seed(seed)
     
     if collage_style != "No Style":
         prompt, negative_prompt = apply_style(collage_style, prompt, negative_prompt)
     else:
         prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator().manual_seed(seed)
-
     if not use_negative_prompt:
-        negative_prompt = ""  # type: ignore
-    negative_prompt += default_negative    
-
-    grid_sizes = {
-        "2x1": (2, 1),
-        "1x2": (1, 2),
-        "2x2": (2, 2),
-        "2x3": (2, 3),
-        "3x2": (3, 2),
-        "1x1": (1, 1)
-    }
+        negative_prompt = None  # type: ignore
     
-    grid_size_x, grid_size_y = grid_sizes.get(grid_size, (2, 2))
-    num_images = grid_size_x * grid_size_y
-
-    options = {
-        "prompt": prompt,
-        "negative_prompt": negative_prompt,
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-        "num_inference_steps": 25,
-        "height": height,
-        "width": width,
-        "num_images_per_prompt": num_images,
-        "output_type": "pil",
-    }
-    
-    torch.cuda.empty_cache()  # Clear GPU memory
-    images = pipe(**options).images
-
-    grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
-
-    for i, img in enumerate(images[:num_images]):
-        grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
-
-    unique_name = str(uuid.uuid4()) + ".png"
-    save_image(grid_img, unique_name)
-    return [unique_name], seed
+    output = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
+        output_type="pil",
+    ).images
+
+    return output
 
 examples = [
-    "Portrait of a beautiful woman in a hat, summer outfit, with freckles on her face, in a close up shot, with sunlight, outdoors, in soft light, with a beach background, looking at the camera, with high resolution photography, in the style of Hasselblad X2D50c --ar 85:128 --v 6.0 --style raw",
-    "Flying food photography with [Two Burgers] as the main subject, Splashes of Toppings and Seasonings, [Rocket Lettuce], [Cheddar Flavored Cheese], [Onion], [Pickles], [Special Sauce], [Sesame Bun], [ sea salt crystals] ::3 Capturing the dynamic splashes of food using high-speed photography , photorealistic, surrealism style, [white background], trending background [clean], Minimalist ::2 [Cuware], [Table], [ Steam], [Smoke], [Vegetable Leaves], [Tomato] ::-0.5 Ad Posters, Pro-Grade Color Grading, Studio Lighting, Rim Lights, [Layered Comps], EOS-1D X Mark III, 500px, Behance, concept art"
+    "A red sofa on top of a white building.",
+    "A cardboard which is large and sits on a theater stage.",
+    "A painting of an astronaut riding a pig wearing a tutu holding a pink umbrella.",
+    "Studio photograph closeup of a chameleon over a black background.",
+    "Closeup portrait photo of beautiful goth woman, makeup.",
+    "A living room, bright modern Scandinavian style house, large windows.",
+    "Portrait photograph of an anthropomorphic tortoise seated on a New York City subway train.",
+    "Batman, cute modern Disney style, Pixar 3d portrait, ultra detailed, gorgeous, 3d zbrush, trending on dribbble, 8k render.",
+    "Cinnamon bun on the plate, watercolor painting, detailed, brush strokes, light palette, light, cozy.",
+    "A lion, colorful, low-poly, cyan and orange eyes, poly-hd, 3d, low-poly game art, polygon mesh, jagged, blocky, wireframe edges, centered composition.",
+    "Long exposure photo of Tokyo street, blurred motion, streaks of light, surreal, dreamy, ghosting effect, highly detailed.",
+    "A glamorous digital magazine photoshoot, a fashionable model wearing avant-garde clothing, set in a futuristic cyberpunk roof-top environment, with a neon-lit city background, intricate high fashion details, backlit by vibrant city glow, Vogue fashion photography.",
+    "Masterpiece, best quality, girl, collarbone, wavy hair, looking at viewer, blurry foreground, upper body, necklace, contemporary, plain pants, intricate, print, pattern, ponytail, freckles, red hair, dappled sunlight, smile, happy."
 ]
 
 css = '''
-.gradio-container{max-width: 560px !important}
+.gradio-container{max-width: 1000px !important}
 h1{text-align:center}
 '''
-with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
-    )
+with gr.Blocks(css=css) as demo:
+    with gr.Row():
+        with gr.Column():
+            gr.HTML(
+            """
+            <h1 style='text-align: center'>
+            Stable Diffusion 3 Medium 
+            </h1>
+            """
+        )
+            gr.HTML(
+                """
+               
+                """
+        )
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
@@ -267,103 +235,77 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            run_button = gr.Button("Run")
-        result = gr.Gallery(label="Grid", columns=1, preview=True)
-
-    with gr.Row(visible=True):
-        collage_style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
-            choices=COLLAGE_STYLE_NAMES,
-            value=DEFAULT_COLLAGE_STYLE_NAME,
-            label="Collage Template",
-        )
-    with gr.Row(visible=True):
-        grid_size_selection = gr.Dropdown(
-            choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
-            value="2x2",
-            label="Grid Size"
-        )
-    with gr.Row(visible=True):
-        style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
+            run_button = gr.Button("Run", scale=0)
+        result = gr.Gallery(label="Result", elem_id="gallery", show_label=False)
+    with gr.Accordion("Advanced options", open=False):
+        with gr.Row():
+            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                value = "deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW",
+                visible=True,
+            )
+        style_selection = gr.Dropdown(
+            label="Style",
             choices=STYLE_NAMES,
             value=DEFAULT_STYLE_NAME,
-            label="Style",
         )
- 
-    with gr.Accordion("Advanced options", open=False):
-        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
-        negative_prompt = gr.Text(
-            label="Negative prompt",
-            max_lines=1,
-            placeholder="Enter a negative prompt",
-            value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
-            visible=True,
+        collage_style_selection = gr.Dropdown(
+            label="Collage Template",
+            choices=COLLAGE_STYLE_NAMES,
+            value=DEFAULT_COLLAGE_STYLE_NAME,
         )
-        with gr.Row():
-            num_inference_steps = gr.Slider(
-                label="Steps",
-                minimum=10,
-                maximum=30,
-                step=1,
-                value=15,
-            )
-        with gr.Row():
-            num_images_per_prompt = gr.Slider(
-                label="Images",
-                minimum=1,
-                maximum=5,
-                step=1,
-                value=2,
-            )
         seed = gr.Slider(
             label="Seed",
             minimum=0,
             maximum=MAX_SEED,
             step=1,
             value=0,
-            visible=True
+        )
+        steps = gr.Slider(
+            label="Steps",
+            minimum=0,
+            maximum=60,
+            step=1,
+            value=30,
+        )
+        number_image = gr.Slider(
+            label="Number of Image",
+            minimum=1,
+            maximum=4,
+            step=1,
+            value=2,
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-        
         with gr.Row(visible=True):
             width = gr.Slider(
                 label="Width",
-                minimum=512,
-                maximum=2048,
-                step=8,
+                minimum=256,
+                maximum=MAX_IMAGE_SIZE,
+                step=32,
                 value=1024,
             )
             height = gr.Slider(
                 label="Height",
-                minimum=512,
-                maximum=2048,
-                step=8,
+                minimum=256,
+                maximum=MAX_IMAGE_SIZE,
+                step=32,
                 value=1024,
             )
-
-
-            
         with gr.Row():
             guidance_scale = gr.Slider(
                 label="Guidance Scale",
                 minimum=0.1,
-                maximum=20.0,
+                maximum=10,
                 step=0.1,
-                value=6,
+                value=7.0,
             )
- 
-
 
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=[result, seed],
+        outputs=[result],
         fn=generate,
         cache_examples=CACHE_EXAMPLES,
     )
@@ -388,36 +330,17 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
             use_negative_prompt,
             style_selection,
             collage_style_selection,
-            grid_size_selection,
             seed,
             width,
             height,
             guidance_scale,
             randomize_seed,
+            steps,
+            number_image,
         ],
-        outputs=[result, seed],
+        outputs=[result],
         api_name="run",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
-
-# Updated inference function
-@spaces.GPU
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image, seed
+    demo.queue().launch()