v1

.gitignore

@@ -0,0 +1,5 @@
+venv
+.git
+.vs
+outputs
+previewer

app.py

@@ -8,29 +8,42 @@ from typing import List
 from diffusers.utils import numpy_to_pil
 from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
-import spaces 
 from previewer.modules import Previewer
-import user_history
+import os
+import datetime
+import json
+import io
+import argparse  # Import the argparse library
+
+# Set up argument parser
+parser = argparse.ArgumentParser(description="Gradio interface for text-to-image generation with optional features.")
+parser.add_argument("--share", action="store_true", help="Enable Gradio sharing.")
+parser.add_argument("--lowvram", action="store_true", help="Enable CPU offload for model operations.")
+parser.add_argument("--torch_compile", action="store_true", help="Enable CPU offload for model operations.")
+
+# Parse arguments
+args = parser.parse_args()
+share = args.share
+ENABLE_CPU_OFFLOAD = args.lowvram  # Use the offload argument to toggle ENABLE_CPU_OFFLOAD
+USE_TORCH_COMPILE = args.torch_compile  # Use the offload argument to toggle ENABLE_CPU_OFFLOAD
 
-os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 
-DESCRIPTION = "# Stable Cascade"
-DESCRIPTION += "\n<p style=\"text-align: center\">Unofficial demo for <a href='https://huggingface.co/stabilityai/stable-cascade' target='_blank'>Stable Casacade</a>, a new high resolution text-to-image model by Stability AI, built on the Würstchen architecture - <a href='https://huggingface.co/stabilityai/stable-cascade/blob/main/LICENSE' target='_blank'>non-commercial research license</a></p>"
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+DESCRIPTION = "<p style=\"font-size:14px\">Stable Cascade Modified By SECourses - Unofficial demo for <a href='https://huggingface.co/stabilityai/stable-cascade' target='_blank'>Stable Casacade</a>, a new high resolution text-to-image model by Stability AI, built on the Würstchen architecture.<br/> Some tips: Higher batch size working great with fast speed and not much VRAM usage - Not all resolutions working e.g. 1920x1080 fails but 1920x1152 works<br/>Supports high resolutions very well such as 1536x1536</p>"
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶</p>"
+    DESCRIPTION += "<br/><p>Running on CPU 🥶</p>"
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") != "0"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1536"))
-USE_TORCH_COMPILE = False
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
+MAX_IMAGE_SIZE = 2048
 PREVIEW_IMAGES = True
 
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-    prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype).to(device)
-    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype).to(device) 
+    prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype)
+    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype)
+    prior_pipeline.enable_xformers_memory_efficient_attention()
+    decoder_pipeline.enable_xformers_memory_efficient_attention()
 
     if ENABLE_CPU_OFFLOAD:
         prior_pipeline.enable_model_cpu_offload()
@@ -66,7 +79,6 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -74,172 +86,159 @@ def generate(
     width: int = 1024,
     height: int = 1024,
     prior_num_inference_steps: int = 30,
-    # prior_timesteps: List[float] = None,
     prior_guidance_scale: float = 4.0,
     decoder_num_inference_steps: int = 12,
-    # decoder_timesteps: List[float] = None,
     decoder_guidance_scale: float = 0.0,
-    num_images_per_prompt: int = 2,
-    profile: gr.OAuthProfile | None = None,
-) -> PIL.Image.Image:
-    #prior_pipeline.to(device)
-    #decoder_pipeline.to(device)
-    #previewer.eval().requires_grad_(False).to(device).to(dtype)
-    generator = torch.Generator().manual_seed(seed)
-    prior_output = prior_pipeline(
-        prompt=prompt,
-        height=height,
-        width=width,
-        num_inference_steps=prior_num_inference_steps,
-        timesteps=DEFAULT_STAGE_C_TIMESTEPS,
-        negative_prompt=negative_prompt,
-        guidance_scale=prior_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-        generator=generator,
-        callback=callback_prior,
-        callback_steps=callback_steps
-    )
-
-    if PREVIEW_IMAGES:
-        for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
-            r = next(prior_output)
-            if isinstance(r, list):
-                yield r[0]
-        prior_output = r
+    batch_size_per_prompt: int = 2,
+    number_of_images_per_prompt: int = 1,  # New parameter
+) -> List[PIL.Image.Image]:
+    images = []  # Initialize an empty list to collect generated images
+    original_seed = seed  # Store the original seed value
+    for i in range(number_of_images_per_prompt):
+        if i > 0:  # Update seed for subsequent iterations
+            seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator().manual_seed(seed)
+        prior_output = prior_pipeline(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=prior_num_inference_steps,
+            timesteps=DEFAULT_STAGE_C_TIMESTEPS,
+            negative_prompt=negative_prompt,
+            guidance_scale=prior_guidance_scale,
+            num_images_per_prompt=batch_size_per_prompt,
+            generator=generator,
+            callback=callback_prior,
+            callback_steps=callback_steps
+        )
 
-    decoder_output = decoder_pipeline(
-        image_embeddings=prior_output.image_embeddings,
-        prompt=prompt,
-        num_inference_steps=decoder_num_inference_steps,
-        # timesteps=decoder_timesteps,
-        guidance_scale=decoder_guidance_scale,
-        negative_prompt=negative_prompt,
-        generator=generator,
-        output_type="pil",
-    ).images
+        if PREVIEW_IMAGES:
+            for _ in range(len(DEFAULT_STAGE_C_TIMESTEPS)):
+                r = next(prior_output)
+            prior_output = r
 
-    #Save images
-    for image in decoder_output:
-        user_history.save_image(
-            profile=profile,
-            image=image,
-            label=prompt,
-            metadata={
-                "negative_prompt": negative_prompt,
-                "seed": seed,
-                "width": width,
-                "height": height,
-                "prior_guidance_scale": prior_guidance_scale,
-                "decoder_num_inference_steps": decoder_num_inference_steps,
-                "decoder_guidance_scale": decoder_guidance_scale,
-                "num_images_per_prompt": num_images_per_prompt,
-            },
-        )
+        decoder_output = decoder_pipeline(
+            image_embeddings=prior_output.image_embeddings,
+            prompt=prompt,
+            num_inference_steps=            decoder_num_inference_steps,
+            guidance_scale=decoder_guidance_scale,
+            negative_prompt=negative_prompt,
+            generator=generator,
+            output_type="pil",
+        ).images
 
-    yield decoder_output[0]
+        # Append generated images to the images list
+        images.extend(decoder_output)
 
+        # Optionally, save each image
+        output_folder = 'outputs'
+        if not os.path.exists(output_folder):
+            os.makedirs(output_folder)
+        for image in decoder_output:
+            # Generate timestamped filename
+            timestamp = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S_%f')
+            image_filename = f"{output_folder}/{timestamp}.png"
+            image.save(image_filename)
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A mecha robot in a favela by Tarsila do Amaral",
-    "The sprirt of a Tamagotchi wandering in the city of Los Angeles",
-    "A delicious feijoada ramen dish"
-]
+    # Return the list of generated images
+    return images
 
-with gr.Blocks() 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.Group():
-        with gr.Row():
+with gr.Blocks() as app:
+    with gr.Row():
+        gr.Markdown(DESCRIPTION)
+    with gr.Row():
+        with gr.Column():
             prompt = gr.Text(
                 label="Prompt",
-                show_label=False,
-                max_lines=1,
                 placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Run", scale=0)
-        result = gr.Image(label="Result", show_label=False)
-    with gr.Accordion("Advanced options", open=False):
-        negative_prompt = gr.Text(
-            label="Negative prompt",
-            max_lines=1,
-            placeholder="Enter a Negative Prompt",
-        )
-
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0,
-        )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row():
-            width = gr.Slider(
-                label="Width",
-                minimum=1024,
-                maximum=MAX_IMAGE_SIZE,
-                step=512,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=1024,
-                maximum=MAX_IMAGE_SIZE,
-                step=512,
-                value=1024,
-            )
-            num_images_per_prompt = gr.Slider(
-                label="Number of Images",
-                minimum=1,
-                maximum=2,
-                step=1,
-                value=1,
             )
-        with gr.Row():
-            prior_guidance_scale = gr.Slider(
-                label="Prior Guidance Scale",
-                minimum=0,
-                maximum=20,
-                step=0.1,
-                value=4.0,
-            )
-            prior_num_inference_steps = gr.Slider(
-                label="Prior Inference Steps",
-                minimum=10,
-                maximum=30,
-                step=1,
-                value=20,
+            run_button = gr.Button("Generate")
+            
+            # Advanced options now directly visible
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                placeholder="Enter a Negative Prompt",
             )
 
-            decoder_guidance_scale = gr.Slider(
-                label="Decoder Guidance Scale",
+            seed = gr.Slider(
+                label="Seed",
                 minimum=0,
-                maximum=0,
-                step=0.1,
-                value=0.0,
-            )
-            decoder_num_inference_steps = gr.Slider(
-                label="Decoder Inference Steps",
-                minimum=4,
-                maximum=12,
+                maximum=MAX_SEED,
                 step=1,
-                value=10,
+                value=0,
             )
-
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=result,
-        fn=generate,
-        cache_examples=CACHE_EXAMPLES,
-    )
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Row():
+                with gr.Column():
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=512,
+                        maximum=MAX_IMAGE_SIZE,
+                        step=64,
+                        value=1024,
+                    )
+                with gr.Column():
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=512,
+                        maximum=MAX_IMAGE_SIZE,
+                        step=64,
+                        value=1024,
+                    )
+            with gr.Row():
+                with gr.Column():
+                    batch_size_per_prompt = gr.Slider(
+                        label="Batch Size",
+                        minimum=1,
+                        maximum=20,
+                        step=1,
+                        value=1,
+                    )
+                with gr.Column():
+                    number_of_images_per_prompt = gr.Slider(
+                        label="Number Of Images To Generate",
+                        minimum=1,
+                        maximum=9999999,
+                        step=1,
+                        value=1,
+                    )
+            with gr.Row():
+                with gr.Column():
+                    prior_guidance_scale = gr.Slider(
+                        label="Prior Guidance Scale (CFG)",
+                        minimum=0,
+                        maximum=20,
+                        step=0.1,
+                        value=4.0,
+                    )
+                with gr.Column():
+                    decoder_guidance_scale = gr.Slider(
+                        label="Decoder Guidance Scale (CFG)",
+                        minimum=0,
+                        maximum=20,
+                        step=0.1,
+                        value=0.0,
+                    )
+            with gr.Row():
+                with gr.Column():
+                    prior_num_inference_steps = gr.Slider(
+                        label="Prior Inference Steps",
+                        minimum=1,
+                        maximum=100,
+                        step=1,
+                        value=20,
+                    )
+                with gr.Column():
+                    decoder_num_inference_steps = gr.Slider(
+                        label="Decoder Inference Steps",
+                        minimum=1,
+                        maximum=100,
+                        step=1,
+                        value=20,
+                    )
+            
+        with gr.Column():
+            result = gr.Gallery(label="Result", show_label=False, height=768)
 
     inputs = [
             prompt,
@@ -253,7 +252,8 @@ with gr.Blocks() as demo:
             decoder_num_inference_steps,
             # decoder_timesteps,
             decoder_guidance_scale,
-            num_images_per_prompt,
+            batch_size_per_prompt,
+            number_of_images_per_prompt
     ]
     gr.on(
         triggers=[prompt.submit, negative_prompt.submit, run_button.click],
@@ -268,12 +268,6 @@ with gr.Blocks() as demo:
         outputs=result,
         api_name="run",
     )
-    
-with gr.Blocks(css="style.css") as demo_with_history:
-    with gr.Tab("App"):
-        demo.render()
-    with gr.Tab("Past generations"):
-        user_history.render()
-
+		
 if __name__ == "__main__":
-    demo_with_history.queue(max_size=20).launch()
+    app.queue().launch(share=share,inbrowser=True)