requirements.txt

gradio>=4.0.0
torch>=2.0.0
diffusers>=0.25.0
transformers>=4.35.0
accelerate>=0.24.0
sentencepiece>=0.1.99
pillow>=10.0.0
numpy>=1.24.0

app.py

@@ -0,0 +1,147 @@
+# main.py
+import gradio as gr
+import numpy as np
+import random
+import torch
+import os
+from diffusers import SanaSprintPipeline
+from PIL import Image
+
+# Initialize device and dtype
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load models
+pipe = SanaSprintPipeline.from_pretrained(
+    "Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers",
+    torch_dtype=dtype
+)
+pipe2 = SanaSprintPipeline.from_pretrained(
+    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
+    torch_dtype=dtype
+)
+pipe.to(device)
+pipe2.to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+def generate_image(prompt, model_size, seed, randomize_seed, width, height, guidance_scale, steps):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    
+    selected_pipe = pipe if model_size == "0.6B" else pipe2
+    
+    result = selected_pipe(
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=steps,
+        width=width,
+        height=height,
+        generator=generator,
+        output_type="pil"
+    )
+    
+    image = result.images[0]
+    filename = f"output_{seed}.png"
+    image.save(filename)
+    return image, filename, seed
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 800px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("# 🚀 Sana Sprint Image Generator")
+        
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(
+                    label="Enter Prompt",
+                    placeholder="A surreal landscape with...",
+                    lines=3
+                )
+                
+                model_size = gr.Radio(
+                    label="Model Size",
+                    choices=["0.6B", "1.6B"],
+                    value="1.6B"
+                )
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        value=42,
+                        step=1
+                    )
+                    randomize_seed = gr.Checkbox(
+                        label="Randomize Seed",
+                        value=True
+                    )
+                    
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            value=1024,
+                            step=32
+                        )
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            value=1024,
+                            step=32
+                        )
+                    
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1.0,
+                        maximum=15.0,
+                        value=4.5,
+                        step=0.1
+                    )
+                    
+                    steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=1,
+                        maximum=50,
+                        value=2,
+                        step=1
+                    )
+                
+                generate_btn = gr.Button("Generate Image", variant="primary")
+
+            with gr.Column():
+                output_image = gr.Image(label="Generated Image")
+                file_output = gr.File(label="Download Image")
+                seed_info = gr.Textbox(label="Used Seed")
+        
+        gr.Examples(
+            examples=[
+                ["a tiny astronaut hatching from an egg on the moon", "1.6B"],
+                ["🐶 Wearing 🕶 flying on the 🌈", "1.6B"],
+                ["an anime illustration of a wiener schnitzel", "0.6B"]
+            ],
+            inputs=[prompt, model_size],
+            outputs=[output_image, file_output, seed_info],
+            fn=generate_image,
+            cache_examples=True
+        )
+
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt, model_size, seed, randomize_seed, width, height, guidance_scale, steps],
+        outputs=[output_image, file_output, seed_info]
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")