# ============================================================ # IMPORTANT: imports order matters for Hugging Face Spaces # ============================================================ import os import gc import random # ---- Spaces GPU decorator (must be imported early) ---------- try: import spaces SPACES_AVAILABLE = True except Exception: SPACES_AVAILABLE = False import gradio as gr import numpy as np from PIL import Image import torch from diffusers import ( StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler, ) from huggingface_hub import login from compel import CompelForSDXL # ============================================================ # Config # ============================================================ MODEL_ID = "telcom/deewaiREALCN" REVISION = "main" #n_680 HF_TOKEN = os.getenv("HF_TOKEN", "").strip() if HF_TOKEN: login(token=HF_TOKEN) MAX_SEED = np.iinfo(np.int32).max # ============================================================ # Device & dtype # ============================================================ cuda_available = torch.cuda.is_available() device = torch.device("cuda" if cuda_available else "cpu") dtype = torch.float16 if cuda_available else torch.float32 MAX_IMAGE_SIZE = 1216 if cuda_available else 768 fallback_msg = "" if not cuda_available: fallback_msg = "GPU unavailable. Running in CPU fallback mode." # ============================================================ # Load pipelines # ============================================================ pipe_txt2img = None pipe_img2img = None compel = None model_loaded = False load_error = None try: from_pretrained_kwargs = { "torch_dtype": dtype, "use_safetensors": True, } if cuda_available: from_pretrained_kwargs["variant"] = "fp16" if HF_TOKEN: from_pretrained_kwargs["token"] = HF_TOKEN pipe_txt2img = StableDiffusionXLPipeline.from_pretrained( MODEL_ID, revision=REVISION, **from_pretrained_kwargs, ) pipe_txt2img.scheduler = EulerAncestralDiscreteScheduler.from_config( pipe_txt2img.scheduler.config ) pipe_txt2img = pipe_txt2img.to(device) # Memory optimisations pipe_txt2img.enable_vae_slicing() pipe_txt2img.enable_attention_slicing() try: pipe_txt2img.enable_xformers_memory_efficient_attention() except Exception: pass pipe_txt2img.set_progress_bar_config(disable=True) # img2img pipeline shares weights pipe_img2img = StableDiffusionXLImg2ImgPipeline(**pipe_txt2img.components) pipe_img2img.scheduler = EulerAncestralDiscreteScheduler.from_config( pipe_img2img.scheduler.config ) pipe_img2img = pipe_img2img.to(device) compel = CompelForSDXL(pipe_txt2img, device=str(device)) model_loaded = True except Exception as e: load_error = repr(e) model_loaded = False # ============================================================ # Utility: error image # ============================================================ def make_error_image(w, h): return Image.new("RGB", (w, h), (18, 18, 22)) # ============================================================ # Inference function # IMPORTANT: decorator must be explicit # ============================================================ if SPACES_AVAILABLE: @spaces.GPU def infer( prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, init_image, strength, ): return _infer_impl( prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, init_image, strength, ) else: def infer( prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, init_image, strength, ): return _infer_impl( prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, init_image, strength, ) def _infer_impl( prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, init_image, strength, ): width = int(width) height = int(height) seed = int(seed) if not model_loaded: return make_error_image(width, height), f"Model load failed: {load_error}" if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) status = f"Seed: {seed}" if fallback_msg: status += f" | {fallback_msg}" try: with torch.inference_mode(): conditioning = compel(prompt, negative_prompt=negative_prompt) common_kwargs = dict( prompt_embeds=conditioning.embeds, pooled_prompt_embeds=conditioning.pooled_embeds, negative_prompt_embeds=conditioning.negative_embeds, negative_pooled_prompt_embeds=conditioning.negative_pooled_embeds, guidance_scale=float(guidance_scale), num_inference_steps=int(num_inference_steps), generator=generator, ) if device.type == "cuda": with torch.autocast("cuda", dtype=dtype): if init_image is not None: image = pipe_img2img( image=init_image, strength=float(strength), **common_kwargs, ).images[0] else: image = pipe_txt2img( width=width, height=height, **common_kwargs, ).images[0] else: if init_image is not None: image = pipe_img2img( image=init_image, strength=float(strength), **common_kwargs, ).images[0] else: image = pipe_txt2img( width=width, height=height, **common_kwargs, ).images[0] return image, status except Exception as e: return make_error_image(width, height), f"Error: {type(e).__name__}: {e}" finally: gc.collect() if device.type == "cuda": torch.cuda.empty_cache() # ============================================================ # UI # ============================================================ CSS = """ body { background: #000; color: #fff; } """ with gr.Blocks(title="SDXL txt2img + img2img") as demo: gr.HTML(f"") if fallback_msg: gr.Markdown(f"**{fallback_msg}**") if not model_loaded: gr.Markdown(f"⚠️ Model failed to load:\n\n{load_error}") gr.Markdown("## SDXL Generator (txt2img + img2img)") prompt = gr.Textbox(label="Prompt", lines=2) init_image = gr.Image(label="Initial image (optional)", type="pil") run_button = gr.Button("Generate") result = gr.Image(label="Result") status = gr.Markdown("") with gr.Accordion("Advanced Settings", open=False): negative_prompt = gr.Textbox(label="Negative prompt") seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed") randomize_seed = gr.Checkbox(value=True, label="Randomize seed") width = gr.Slider(256, MAX_IMAGE_SIZE, step=32, value=512, label="Width") height = gr.Slider(256, MAX_IMAGE_SIZE, step=32, value=512, label="Height") guidance_scale = gr.Slider(0, 20, step=0.1, value=7, label="Guidance scale") num_inference_steps = gr.Slider(1, 40, step=1, value=20, label="Steps") strength = gr.Slider(0.0, 1.0, step=0.05, value=0.7, label="Image strength") run_button.click( fn=infer, inputs=[ prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, init_image, strength, ], outputs=[result, status], ) demo.queue().launch(ssr_mode=False)