import os import uuid import random from typing import Tuple, Optional import gradio as gr import numpy as np from PIL import Image import torch import spaces from diffusers import ( StableDiffusionXLPipeline, StableDiffusionPipeline, EulerAncestralDiscreteScheduler, ) PRIMARY_MODEL_ID = "SG161222/RealVisXL_V5.0_Lightning" # requires access + token FALLBACK_MODEL_ID = "stabilityai/sd-turbo" # public, fast 1.5-turbo def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]: styles = { "3840 x 2160": ( "hyper-realistic image of {prompt}. lifelike, authentic, natural colors, " "true-to-life details, landscape image, realistic lighting, immersive, highly detailed", "unrealistic, low resolution, artificial, over-saturated, distorted, fake", ), "Style Zero": ("{prompt}", ""), } DEFAULT_STYLE_NAME = "3840 x 2160" p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) n2 = (n + (" " + negative if negative else "")).strip() return p.replace("{prompt}", positive), n2 def _enable_performance_knobs(): if torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True torch.set_grad_enabled(False) def _try_load_realvis(hf_token: Optional[str]): use_cuda = torch.cuda.is_available() dtype = torch.float16 if use_cuda else torch.float32 pipe = StableDiffusionXLPipeline.from_pretrained( PRIMARY_MODEL_ID, torch_dtype=dtype, use_safetensors=True, add_watermarker=False, token=hf_token, # <- IMPORTANT ) pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) try: pipe.enable_xformers_memory_efficient_attention() except Exception: pass device = torch.device("cuda:0" if use_cuda else "cpu") pipe = pipe.to(device) return pipe def _try_load_fallback(): # sd-turbo is Stable Diffusion 1.5 turbo; quick & public use_cuda = torch.cuda.is_available() dtype = torch.float16 if use_cuda else torch.float32 pipe = StableDiffusionPipeline.from_pretrained( FALLBACK_MODEL_ID, torch_dtype=dtype, use_safetensors=True, ) try: pipe.enable_xformers_memory_efficient_attention() except Exception: pass device = torch.device("cuda:0" if use_cuda else "cpu") pipe = pipe.to(device) return pipe def load_and_prepare_model(): _enable_performance_knobs() hf_token = os.getenv("HF_TOKEN", "").strip() or None # Try RealVis first try: return _try_load_realvis(hf_token) except Exception as e: msg = str(e).lower() if ("401" in msg or "403" in msg or "unauthorized" in msg or "forbidden" in msg): # Clear hint in server logs; UI will still work via fallback. print( "\n[WARNING] Could not load RealVisXL (auth). " "Make sure you've requested access and set HF_TOKEN in Space secrets.\n" ) else: print(f"\n[WARNING] RealVisXL failed to load: {e}\n") # Fallback to sd-turbo so app still runs print("[INFO] Falling back to stabilityai/sd-turbo (public).") return _try_load_fallback() # Load once model = load_and_prepare_model() def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: if randomize_seed: seed = random.randint(0, np.iinfo(np.int32).max) return int(seed) def save_image(img: Image.Image) -> str: unique_name = f"{uuid.uuid4().hex}.png" img.save(unique_name) return unique_name @spaces.GPU(duration=60, enable_queue=True) def generate( prompt: str, seed: int = 1, width: int = 1024, height: int = 1024, guidance_scale: float = 3.0, num_inference_steps: int = 25, randomize_seed: bool = False, ): if not prompt or not prompt.strip(): raise gr.Error("Please enter a prompt.") # Make dimensions friendly for SD models width = int(max(256, (width // 8) * 8)) height = int(max(256, (height // 8) * 8)) seed = randomize_seed_fn(seed, randomize_seed) generator = torch.Generator(device=model.device).manual_seed(seed) # If model is SDXL pipeline, use the styled prompts; if fallback SD1.5 turbo, style still OK positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt) # Some pipelines (sd-turbo) ignore guidance/steps or behave differently; passing is still safe out = model( prompt=positive_prompt, negative_prompt=negative_prompt, width=width if "xl" in model.__class__.__name__.lower() else None, height=height if "xl" in model.__class__.__name__.lower() else None, guidance_scale=float(guidance_scale), num_inference_steps=int(num_inference_steps), generator=generator, output_type="pil", ) # Handle both diffusers return shapes images = getattr(out, "images", out) image_path = save_image(images[0]) return image_path with gr.Blocks(theme="soft") as demo: with gr.Row(): with gr.Column(scale=12, elem_id="title_block"): gr.Markdown( "

SNAPSCRIBE

" ) gr.Markdown( "

Developed with ❤ by Aklavya

" ) with gr.Row(): with gr.Column(scale=3): prompt = gr.Textbox( label="Input Prompt", placeholder="Describe the image you want to create", lines=2, ) seed = gr.Number(value=1, label="Seed", precision=0) randomize_seed = gr.Checkbox(value=True, label="Randomize Seed") width = gr.Slider(512, 1536, value=1024, step=8, label="Width") height = gr.Slider(512, 1536, value=1024, step=8, label="Height") guidance_scale = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Guidance Scale") steps = gr.Slider(10, 35, value=25, step=1, label="Inference Steps") run_button = gr.Button("Generate Image", variant="primary") example_prompts_text = ( "Dew-covered spider web in morning sunlight, with blurred greenery\n" "--------------------------------------------\n" "Glass of cold water with ice cubes and condensation on a wooden table\n" "--------------------------------------------\n" "Coffee cup with latte art, steam rising, and morning sunlight\n" "--------------------------------------------\n" "Autumn forest with golden leaves, sunlight through trees, and a breeze" ) gr.Textbox( value=example_prompts_text, lines=8, label="Sample Inputs", interactive=False, ) with gr.Column(scale=7): result_image = gr.Image( label="Generated Image", type="filepath", elem_id="output_image", ) run_button.click( fn=generate, inputs=[prompt, seed, width, height, guidance_scale, steps, randomize_seed], outputs=[result_image], api_name="generate", ) if __name__ == "__main__": demo.launch()