Deploy Gradio app with multiple files

app.py

@@ -0,0 +1,137 @@
+import gradio as gr
+import torch
+import spaces
+import os
+from diffusers import DiffusionPipeline
+
+# --- Model Configuration and Loading ---
+MODEL_ID = "Manojb/stable-diffusion-2-1-base"
+DTYPE = torch.bfloat16
+
+try:
+    # Load pipeline
+    pipe = DiffusionPipeline.from_pretrained(
+        MODEL_ID, 
+        torch_dtype=DTYPE, 
+        use_safetensors=True
+    )
+    pipe.to('cuda')
+    
+    # --- Mandatory ZeroGPU AoT Compilation for Optimization ---
+    
+    @spaces.GPU(duration=1500)  # Extended duration for startup compilation
+    def compile_unet():
+        print("Starting AoT compilation for UNet...")
+        
+        # Dummy inputs for 512x512 generation (B=1, latents=64x64 for UNet)
+        B, C, H, W = 1, 4, 64, 64
+        sample = torch.randn(B, C, H, W, dtype=DTYPE, device='cuda')
+        timestep = torch.tensor([999], dtype=torch.long, device='cuda')
+        
+        # Encoder Hidden States (text embeddings): (B, 77, 1024) for SD2.1
+        EHS_DIM = 77
+        EHS_HIDDEN = 1024
+        encoder_hidden_states = torch.randn(B, EHS_DIM, EHS_HIDDEN, dtype=DTYPE, device='cuda')
+
+        inputs = (sample, timestep, encoder_hidden_states)
+        
+        with spaces.aoti_capture(pipe.unet) as call:
+            call(*inputs)
+        
+        exported = torch.export.export(pipe.unet, args=call.args, kwargs=call.kwargs)
+        compiled_model = spaces.aoti_compile(exported)
+        print("AoT compilation successful.")
+        return compiled_model
+
+    # Execute compilation during startup
+    compiled_unet = compile_unet()
+    spaces.aoti_apply(compiled_unet, pipe.unet)
+
+except Exception as e:
+    print(f"⚠️ Warning: Model initialization or AoT compilation failed ({e}). Running without optimization or skipping initialization if severe.")
+    # Fallback to loading the model without AoT if compilation fails
+    if 'pipe' not in locals():
+        pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=DTYPE, use_safetensors=True)
+        pipe.to('cuda')
+        print("Model loaded successfully without AoT.")
+
+@spaces.GPU(duration=60) # Standard GPU allocation for inference
+def generate(prompt: str, num_images: int):
+    """Generates images using the Stable Diffusion pipeline."""
+    
+    if not prompt:
+        raise gr.Error("Prompt cannot be empty.")
+
+    # Prepare batch input
+    prompt_list = [prompt] * num_images
+    
+    # Generate images
+    output = pipe(
+        prompt_list,
+        num_inference_steps=25,
+        guidance_scale=9.0,
+    )
+    
+    return output.images
+
+# --- Gradio Interface ---
+
+with gr.Blocks(theme=gr.themes.Soft(), title="SD 2.1 Base Generator") as demo:
+    gr.HTML(
+        """
+        <div style="text-align: center; margin-bottom: 20px;">
+            <h1>Stable Diffusion 2.1 Base (512x512)</h1>
+            <p>Model: Manojb/stable-diffusion-2-1-base | Optimized with ZeroGPU AoT</p>
+            <p>Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">anycoder</a></p>
+        </div>
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="A detailed digital painting of a majestic dragon flying over a medieval castle, fantasy art",
+                lines=3
+            )
+            num_images = gr.Slider(
+                minimum=1,
+                maximum=4,
+                step=1,
+                value=2,
+                label="Number of Images to Generate (Max 4)",
+                info="Generates multiple images in a single batch call."
+            )
+            generate_btn = gr.Button("Generate Images", variant="primary")
+        
+        with gr.Column(scale=2):
+            output_gallery = gr.Gallery(
+                label="Generated Images (512x512)",
+                height=512,
+                columns=2,
+                rows=2,
+                object_fit="contain"
+            )
+
+    generate_btn.click(
+        fn=generate,
+        inputs=[prompt, num_images],
+        outputs=output_gallery
+    )
+
+    gr.Examples(
+        examples=[
+            ["A photorealistic portrait of a golden retriever wearing sunglasses on a beach, cinematic lighting", 2],
+            ["Steampunk owl on a bookshelf, detailed brass gears, oil painting", 4],
+            ["High contrast black and white photograph of an old lighthouse during a storm", 1]
+        ],
+        inputs=[prompt, num_images],
+        outputs=output_gallery,
+        fn=generate,
+        cache_examples=True,
+        cache_mode="eager"
+    )
+
+demo.queue()
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,19 @@
+gradio
+torch
+diffusers
+accelerate
+safetensors
+git+https://github.com/huggingface/spaces
+Pillow
+xformers
+scipy
+opencv-python
+ftfy
+transformers
+regex
+httpx
+pydantic
+typing-extensions
+dataclasses_json
+aiohttp
+numpy