Update app.py

app.py

@@ -1,62 +1,434 @@
+# -*- coding: utf-8 -*-
+"""
+Hugging Face Spaces Script for Basic Stable Diffusion 1.5 Gradio App
+Adapted from user's local script and HF Spaces template.
+Supports Hub models and CPU/GPU selection based on available hardware.
+"""
+
 import gradio as gr
 import numpy as np
 import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 import torch
+from diffusers import StableDiffusionPipeline
+# Import commonly used schedulers
+from diffusers import DDPMScheduler, EulerDiscreteScheduler, DPMSolverMultistepScheduler, LMSDiscreteScheduler
+import os # Keep os for potential checks, though local paths are less standard on Spaces
+# import spaces #[uncomment to use ZeroGPU if needed, typically not for standard GPU usage]
+from PIL import Image
+import time # Optional: for timing generation
+# huggingface_hub is implicitly used by from_pretrained
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+# --- Configuration ---
+# MODELS_DIR is less relevant for Spaces unless specifically placing models in repo
+# For Spaces, models are primarily loaded via their Hugging Face Hub IDs
+SUPPORTED_SD15_SIZES = ["512x512", "768x512", "512x768", "768x768", "1024x768", "768x1024", "1024x1024", "hire.fix"]
 
+# Mapping of friendly scheduler names to their diffusers classes
+SCHEDULER_MAP = {
+    "Euler": EulerDiscreteScheduler,
+    "DPM++ 2M": DPMSolverMultistepScheduler,
+    "DDPM": DDPMScheduler,
+    "LMS": LMSDiscreteScheduler,
+    # Add more as needed from diffusers.schedulers (make sure they are imported)
+}
+DEFAULT_SCHEDULER = "Euler" # Default scheduler on startup
+
+# List of popular Stable Diffusion 1.5 models on the Hugging Face Hub
+# For Spaces, this is the primary source of models.
+DEFAULT_HUB_MODELS = [
+    "runwayml/stable-diffusion-v1-5",
+    "SG161222/Realistic_Vision_V6.0_B1_noVAE", # Example popular 1.5 model
+    "nitrosocke/Ghibli-Diffusion",
+    "danyloylo/sd1.5-ghibli-style-05",
+    "Bilal326/SD_1.5_DragonWarriorV2"
+    # "CompVis/stable-diffusion-v1-4", # Example SD 1.4 model (might behave slightly differently)
+    # Add other diffusers-compatible SD1.5 models here
+]
+
+# --- Determine available devices and set up options ---
+# This logic is from the user's script and works well for Spaces
+AVAILABLE_DEVICES = ["CPU"]
 if torch.cuda.is_available():
-    torch_dtype = torch.float16
+    AVAILABLE_DEVICES.append("GPU")
+    print(f"CUDA available. Found {torch.cuda.device_count()} GPU(s).")
+    if torch.cuda.device_count() > 0:
+        print(f"Using GPU 0: {torch.cuda.get_device_name(0)}")
+else:
+    print("CUDA not available. Running on CPU.")
+
+# Default device preference: GPU if available, else CPU
+DEFAULT_DEVICE = "GPU" if "GPU" in AVAILABLE_DEVICES else "CPU"
+
+# Set initial PyTorch device string based on detection
+initial_device_to_use = "cuda" if DEFAULT_DEVICE == "GPU" else "cpu"
+print(f"Initial pipeline will load on device: {initial_device_to_use}")
+
+# Determine initial dtype
+# Note: fp16 is generally faster and uses less VRAM on compatible GPUs
+initial_dtype_to_use = torch.float32 # Default
+if initial_device_to_use == "cuda":
+    # Check if the GPU supports fp16 (most modern ones do)
+    if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 7: # Check compute capability (7.0+ for good fp16)
+         initial_dtype_to_use = torch.float16
+         print("Detected GPU supports FP16, will attempt using torch.float16.")
+    else:
+         initial_dtype_to_use = torch.float32 # Fallback
+         print("Detected GPU might not fully support FP16 or capability check failed, using torch.float32.")
 else:
-    torch_dtype = torch.float32
+    initial_dtype_to_use = torch.float32 # CPU requires float32
+
+print(f"Initial dtype: {initial_dtype_to_use}")
+
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# --- Global state for the loaded pipeline ---
+# We'll load the *initial* pipeline once on startup and keep it in memory.
+# Subsequent model changes will reload the pipeline within the infer function.
+current_pipeline = None
+current_model_id = None # Keep track of the currently loaded model identifier
+current_device_loaded = None # Keep track of the device the pipeline is currently on
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+# Initial model to load on startup
+INITIAL_MODEL_ID = DEFAULT_HUB_MODELS[0] if DEFAULT_HUB_MODELS else None
+if INITIAL_MODEL_ID:
+    print(f"\nLoading initial model '{INITIAL_MODEL_ID}' on startup...")
+    try:
+        current_pipeline = StableDiffusionPipeline.from_pretrained(
+            INITIAL_MODEL_ID,
+            torch_dtype=initial_dtype_to_use,
+            safety_checker=None, # <<< SAFETY CHECKER DISABLED <<<
+        )
+        current_pipeline = current_pipeline.to(initial_device_to_use)
+        current_model_id = INITIAL_MODEL_ID
+        current_device_loaded = torch.device(initial_device_to_use)
+        print(f"Initial model loaded successfully on {current_device_loaded}.")
+
+         # Basic check for SD1.x architecture
+        unet_config = getattr(current_pipeline, 'unet', None)
+        if unet_config and hasattr(unet_config, 'config') and hasattr(unet_config.config, 'cross_attention_dim'):
+             cross_attn_dim = unet_config.config.cross_attention_dim
+             if cross_attn_dim != 768:
+                 warning_msg = (f"Warning: Loaded model '{INITIAL_MODEL_ID}' might not be a standard SD 1.x model "
+                                f"(expected UNet cross_attention_dim 768, found {cross_attn_dim}). "
+                                "Results may be unexpected.")
+                 print(warning_msg)
+                 # gr.Warning(warning_msg) # Cannot raise Gradio error/warning during startup load
+             else:
+                 print("UNet cross_attention_dim is 768, consistent with SD 1.x.")
+        else:
+             print("Could not check UNet cross_attention_dim for initial model.")
+
+
+    except Exception as e:
+        current_pipeline = None
+        current_model_id = None
+        current_device_loaded = None
+        print(f"Error loading initial model '{INITIAL_MODEL_ID}': {e}")
+        print("Application will start, but image generation may fail if the initial model cannot be loaded.")
+        # Cannot raise gr.Error here as Gradio not fully initialized
+else:
+     print("\nNo default Hub models defined. Application will start without a loaded model.")
+     print("Please select a model from the dropdown to enable generation.")
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
+# --- Image Generation Function (Adapted for Hugging Face Spaces 'infer' signature) ---
+# @spaces.GPU #[uncomment if using ZeroGPU, otherwise standard torch device handles it]
 def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
+    model_identifier,         # From model_dropdown
+    selected_device_str,      # From device_dropdown
+    prompt,                   # From prompt_input
+    negative_prompt,          # From negative_prompt_input
+    steps,                    # From steps_slider
+    cfg_scale,                # From cfg_slider
+    scheduler_name,           # From scheduler_dropdown
+    size,                     # From size_dropdown
+    seed,                     # From seed_input
+    randomize_seed,           # From randomize_seed_checkbox
+    progress=gr.Progress(track_tqdm=True), # Added progress argument from template
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    """Generates an image using the selected model and parameters on the chosen device."""
+    global current_pipeline, current_model_id, current_device_loaded, SCHEDULER_MAP
 
-    generator = torch.Generator().manual_seed(seed)
+    # Check if initial load failed
+    if current_pipeline is None and model_identifier != INITIAL_MODEL_ID:
+         # Try loading the selected model if initial load failed or model is different
+         pass # Logic below handles loading
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    if not model_identifier or model_identifier == "No models found":
+        raise gr.Error(f"No model selected or available. Please select a model from the list.")
+    if not prompt:
+        raise gr.Error("Please enter a prompt.")
 
-    return image, seed
+    # Map selected device string to PyTorch device string
+    device_to_use = "cuda" if selected_device_str == "GPU" and "GPU" in AVAILABLE_DEVICES else "cpu"
+    # If GPU was selected but not available, raise an error specific to this condition
+    if selected_device_str == "GPU" and device_to_use == "cpu":
+         raise gr.Error("GPU selected but CUDA is not available to PyTorch on this Space. Please select CPU or ensure the Space is configured with a GPU and the CUDA version of PyTorch is installed.")
 
+    # Determine dtype based on the actual device being used
+    dtype_to_use = torch.float32 # Default
+    if device_to_use == "cuda":
+        if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 7:
+             dtype_to_use = torch.float16
+        else:
+             dtype_to_use = torch.float32
+    else:
+         dtype_to_use = torch.float32
+
+
+    print(f"Attempting generation on device: {device_to_use}, using dtype: {dtype_to_use}")
+
+    # 1. Load/Switch Model if necessary
+    # Check if the requested model identifier OR the requested device has changed
+    # Use string comparison for current_device_loaded as it's a torch.device object
+    if current_pipeline is None or current_model_id != model_identifier or (current_device_loaded is not None and str(current_device_loaded) != device_to_use):
+        print(f"Loading model: {model_identifier} onto {device_to_use}...")
+        # Clear previous pipeline to potentially free memory *before* loading the new one
+        if current_pipeline is not None:
+             print(f"Unloading previous model '{current_model_id}' from {current_device_loaded}...")
+             # Move pipeline to CPU before deleting if it was on GPU, might help with freeing VRAM
+             if str(current_device_loaded) == "cuda":
+                  try:
+                      current_pipeline.to("cpu")
+                      print("Moved previous pipeline to CPU.")
+                  except Exception as move_e:
+                      print(f"Warning: Failed to move previous pipeline to CPU: {move_e}")
+
+             del current_pipeline
+             current_pipeline = None # Set to None immediately
+             # Attempt to clear CUDA cache if using GPU (from the previous device)
+             if str(current_device_loaded) == "cuda":
+                 try:
+                     torch.cuda.empty_cache()
+                     print("Cleared CUDA cache.")
+                 except Exception as cache_e:
+                     print(f"Warning: Error clearing CUDA cache: {cache_e}")
+
+        # Ensure the device is actually available if not CPU (redundant with earlier check but safe)
+        if device_to_use == "cuda":
+             if not torch.cuda.is_available():
+                  raise gr.Error("CUDA selected but not available. Please select CPU.")
+
+        try:
+            # For Spaces, assume model_identifier is a Hub ID or a path *within the repo*
+            # from_pretrained can handle both.
+            print(f"Attempting to load model from: {model_identifier}")
+            pipeline = StableDiffusionPipeline.from_pretrained(
+                model_identifier,
+                torch_dtype=dtype_to_use,
+                safety_checker=None, # <<< SAFETY CHECKER DISABLED <<<
+                # Add `vae=AutoencoderKL.from_pretrained(...)` if needed for specific models
+            )
+
+            pipeline = pipeline.to(device_to_use) # Move to the selected device
+
+            current_pipeline = pipeline
+            current_model_id = model_identifier
+            current_device_loaded = torch.device(device_to_use)
+
+            # Basic check for SD1.x architecture (cross_attention_dim = 768)
+            unet_config = getattr(pipeline, 'unet', None)
+            if unet_config and hasattr(unet_config, 'config') and hasattr(unet_config.config, 'cross_attention_dim'):
+                 cross_attn_dim = unet_config.config.cross_attention_dim
+                 if cross_attn_dim != 768:
+                     warning_msg = (f"Warning: Loaded model '{model_identifier}' might not be a standard SD 1.x model "
+                                    f"(expected UNet cross_attention_dim 768, found {cross_attn_dim}). "
+                                    "Results may be unexpected or generation might fail.")
+                     print(warning_msg)
+                     gr.Warning(warning_msg)
+                 else:
+                     print("UNet cross_attention_dim is 768, consistent with SD 1.x.")
+            else:
+                 print("Could not check UNet cross_attention_dim.")
+
+
+            print(f"Model '{model_identifier}' loaded successfully on {current_device_loaded} with dtype {dtype_to_use}.")
+
+        except Exception as e:
+            # Reset global state on load failure
+            current_pipeline = None
+            current_model_id = None
+            current_device_loaded = None
+            print(f"Error loading model '{model_identifier}': {e}")
+            error_message_lower = str(e).lower()
+            # Provide more specific error messages based on common exceptions
+            if "cannot find requested files" in error_message_lower or "404 client error" in error_message_lower or "no such file or directory" in error_message_lower:
+                 raise gr.Error(f"Model '{model_identifier}' not found on Hugging Face Hub or in repo files. Check ID/path or internet connection. Error: {e}")
+            elif "checkpointsnotfounderror" in error_message_lower or "valueerror: could not find a valid model structure" in error_message_lower:
+                 raise gr.Error(f"No valid diffusers model at '{model_identifier}'. Ensure it's a diffusers format ID/path. Error: {e}")
+            elif "out of memory" in error_message_lower:
+                 raise gr.Error(f"Out of Memory (OOM) loading model. This Space might not have enough RAM/VRAM for this model. Try a lighter model or select CPU (if available). Error: {e}")
+            elif "cusolver64" in error_message_lower or "cuda driver version" in error_message_lower or "cuda error" in error_message_lower:
+                 raise gr.Error(f"CUDA/GPU Driver/Installation Error on Space: {e}. Check Space hardware or select CPU.")
+            elif "safetensors_rust.safetensorserror" in error_message_lower or "oserror: cannot load" in error_message_lower or "filenotfounderror" in error_message_lower:
+                 raise gr.Error(f"Model file error for '{model_identifier}': {e}. Files might be corrupt or incomplete on the Hub/in repo.")
+            elif "could not import" in error_message_lower or "module not found" in error_message_lower:
+                 raise gr.Error(f"Dependency error: {e}. Ensure required libraries are in requirements.txt.")
+            else:
+                raise gr.Error(f"Failed to load model '{model_identifier}': {e}")
+
+    # Check if pipeline is successfully loaded before proceeding
+    if current_pipeline is None:
+         raise gr.Error("Model failed to load. Cannot generate image.")
+
+
+    # 2. Configure Scheduler
+    selected_scheduler_class = SCHEDULER_MAP.get(scheduler_name)
+    if selected_scheduler_class is None:
+         print(f"Warning: Unknown scheduler '{scheduler_name}'. Using default: {DEFAULT_SCHEDULER}.")
+         selected_scheduler_class = SCHEDULER_MAP[DEFAULT_SCHEDULER]
+         gr.Warning(f"Unknown scheduler '{scheduler_name}'. Using default: {DEFAULT_SCHEDULER}.")
+
+    # Recreate scheduler from config to ensure compatibility with the loaded pipeline
+    try:
+        scheduler_config = current_pipeline.scheduler.config
+        current_pipeline.scheduler = selected_scheduler_class.from_config(scheduler_config)
+        print(f"Scheduler set to: {scheduler_name}")
+    except Exception as e:
+        print(f"Error setting scheduler '{scheduler_name}': {e}")
+        # Attempt to fallback to a default if setting fails
+        try:
+             print(f"Attempting to fallback to default scheduler: {DEFAULT_SCHEDULER}")
+             current_pipeline.scheduler = SCHEDULER_MAP[DEFAULT_SCHEDULER].from_config(scheduler_config)
+             gr.Warning(f"Failed to set scheduler to '{scheduler_name}', fell back to {DEFAULT_SCHEDULER}. Error: {e}")
+        except Exception as fallback_e:
+             print(f"Fallback scheduler failed too: {fallback_e}")
+             raise gr.Error(f"Failed to configure scheduler '{scheduler_name}' and fallback failed. Error: {e}")
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
 
+    # 3. Parse Image Size
+    width, height = 512, 512 # Default size
+    if size.lower() == "hire.fix":
+        width, height = 1024, 1024
+        print(f"Interpreting 'hire.fix' size as {width}x{height}")
+    else:
+        try:
+            w_str, h_str = size.split('x')
+            width = int(w_str)
+            height = int(h_str)
+        except ValueError:
+            raise gr.Error(f"Invalid size format: '{size}'. Use 'WidthxHeight' (e.g., 512x512) or 'hire.fix'.")
+        except Exception as e:
+             raise gr.Error(f"Error parsing size '{size}': {e}")
+
+    # Size multiple check (SD 1.5 works best with multiples of 64 or 8)
+    multiple_check = 64 # Use 64 as a standard check
+    if width % multiple_check != 0 or height % multiple_check != 0:
+         warning_msg_size = (f"Warning: Image size {width}x{height} is not a multiple of {multiple_check}. "
+                             f"Stable Diffusion 1.5 models are typically trained on sizes like 512x512. "
+                             "Using non-standard sizes may cause tiling, distortions, or other artifacts.")
+         print(warning_msg_size)
+         gr.Warning(warning_msg_size)
+         # Optional: Round size to nearest multiple of 64? Not implemented here to preserve user choice.
+
+
+    # 4. Set Seed Generator
+    generator = None
+    # The generator device needs to match the pipeline device
+    generator_device = current_pipeline.device # Must match the pipeline device
+
+    if randomize_seed: # Use the randomize_seed checkbox
+        seed = random.randint(0, MAX_SEED) # Re-randomize seed
+        print(f"Randomizing seed to: {seed}")
+    else:
+        # Use provided seed if not randomizing (-1 will still use it)
+        print(f"Using provided seed: {int(seed)}")
+
+
+    try:
+        # Explicitly move generator to the desired device
+        generator = torch.Generator(device=generator_device).manual_seed(int(seed))
+        print(f"Generator set with seed {int(seed)} on device: {generator_device}")
+    except Exception as e:
+         print(f"Warning: Error setting seed generator on device {generator_device}: {e}. Generation might still proceed with a default generator (potentially on CPU).")
+         gr.Warning(f"Failed to set seed generator on device {generator_device}. Generation might use a random seed on a different device. Error: {e}")
+         generator = None # Let pipeline handle random seed if generator creation fails or device mismatch
+
+    # 5. Generate Image
+    # Ensure required parameters are integers/floats
+    num_inference_steps_int = int(steps)
+    guidance_scale_float = float(cfg_scale)
+
+    # Basic validation on parameters
+    if num_inference_steps_int <= 0 or guidance_scale_float <= 0:
+         raise ValueError("Steps and CFG Scale must be positive values.")
+    if width <= 0 or height <= 0:
+         raise ValueError("Image width and height must be positive.")
+
+    print(f"Generating: Prompt='{prompt[:80]}{'...' if len(prompt) > 80 else ''}', NegPrompt='{negative_prompt[:80]}{'...' if len(negative_prompt) > 80 else ''}', Steps={num_inference_steps_int}, CFG={guidance_scale_float}, Size={width}x{height}, Scheduler={scheduler_name}, Seed={int(seed)}, Device={device_to_use}, Dtype={dtype_to_use}")
+    start_time = time.time()
+
+    try:
+        # Use the progress parameter from the template
+        output = current_pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt if negative_prompt else None,
+            num_inference_steps=num_inference_steps_int,
+            guidance_scale=guidance_scale_float,
+            width=width,
+            height=height,
+            generator=generator,
+            # Pass progress object for tqdm tracking in Gradio
+            callback_steps=max(1, num_inference_steps_int // 20), # Update progress bar periodically
+            callback=lambda step, timestep, latents: progress((step / num_inference_steps_int, f"Step {step}/{num_inference_steps_int}")),
+
+            # Add VAE usage here if needed for specific models that require it
+            # vae=...
+            # Potentially add attention slicing/xformers/etc. for memory efficiency
+            # enable_attention_slicing="auto", # Can help with VRAM on smaller GPUs
+            # enable_xformers_memory_efficient_attention() # Needs xformers installed & compatible GPU
+        )
+        end_time = time.time()
+        print(f"Generation finished in {end_time - start_time:.2f} seconds.")
+        generated_image = output.images[0]
+
+        # Return both the image and the seed (potentially randomized)
+        return generated_image, seed
+
+    except gr.Error as e:
+         # Re-raise Gradio errors directly
+         raise e
+    except ValueError as ve:
+         # Handle specific value errors like invalid parameters
+         print(f"Parameter Error: {ve}")
+         raise gr.Error(f"Invalid Parameter: {ve}")
+    except Exception as e:
+        # Catch any other unexpected errors during generation
+        print(f"An error occurred during image generation: {e}")
+        error_message_lower = str(e).lower()
+        if "size must be a multiple of" in error_message_lower or "invalid dimensions" in error_message_lower or "shape mismatch" in error_message_lower:
+             raise gr.Error(f"Image generation failed - Invalid size '{width}x{height}' for model: {e}. Try a multiple of 64 or 8.")
+        elif "out of memory" in error_message_lower or "cuda out of memory" in error_message_lower:
+             print("Hint: Try smaller image size, fewer steps, or a model that uses less VRAM.")
+             raise gr.Error(f"Out of Memory (OOM) during generation. This Space might not have enough VRAM. Try smaller size/steps or select CPU (if available). Error: {e}")
+        elif "runtimeerror" in error_message_lower:
+             raise gr.Error(f"Runtime Error during generation: {e}. This could be a model/scheduler incompatibility or other issue.")
+        elif "device-side assert" in error_message_lower or "cuda error" in error_message_lower:
+             raise gr.Error(f"CUDA/GPU Error during generation: {e}. Ensure the Space is configured with a GPU and compatible PyTorch.")
+        elif "expected all tensors to be on the same device" in error_message_lower:
+             raise gr.Error(f"Device mismatch error during generation: {e}. This is an internal error, please report it.")
+        else:
+             # Generic catch-all for unknown errors
+             raise gr.Error(f"Image generation failed: An unexpected error occurred. {e}")
+
+
+# --- Gradio Interface ---
+# For Spaces, we primarily list Hub models in the dropdown
+model_choices = DEFAULT_HUB_MODELS
+if not model_choices:
+    initial_model_choices = ["No models found"]
+    initial_default_model = "No models found"
+    model_dropdown_interactive = False
+    print(f"\n!!! WARNING: No default Hub models listed in script. Model dropdown will be empty. !!!")
+else:
+    initial_model_choices = model_choices
+    # Set a reasonable default if available
+    initial_default_model = INITIAL_MODEL_ID if INITIAL_MODEL_ID else "No models found"
+    model_dropdown_interactive = True if INITIAL_MODEL_ID else False
+
+
+scheduler_choices = list(SCHEDULER_MAP.keys())
+
+# Use the template's CSS
 css = """
 #col-container {
     margin: 0 auto;
@@ -64,91 +436,156 @@ css = """
 }
 """
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: # Added Soft theme from user's script
+    gr.Markdown(
+        """
+        # CipherCore Stable Diffusion 1.5 Generator
+        Create images with Stable Diffusion 1.5 using models from Hugging Face Hub.
+        Choose a model, set your prompts and parameters, and generate!
+        _Note: 'hire.fix' size option currently generates at 1024x1024._
+        """ # Removed reference to local checkpoints and "coming soon"
+    )
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+    # Add a note about model loading time
+    if INITIAL_MODEL_ID:
+         gr.Markdown(f"*(Note: The initial model '{INITIAL_MODEL_ID}' is loading... First generation might take longer.)*")
+    else:
+         gr.Markdown(f"*(Note: No initial model configured or loaded. Select a model from the dropdown to start.)*")
 
-        result = gr.Image(label="Result", show_label=False)
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+    with gr.Row():
+        with gr.Column(scale=2): # Give more space to controls
+            model_dropdown = gr.Dropdown(
+                choices=initial_model_choices,
+                value=initial_default_model,
+                label="Select Model (Hugging Face Hub ID)", # Updated label
+                interactive=model_dropdown_interactive,
             )
+            device_dropdown = gr.Dropdown(
+                choices=AVAILABLE_DEVICES,
+                value=DEFAULT_DEVICE,
+                label="Processing Device",
+                interactive=len(AVAILABLE_DEVICES) > 1, # Only make interactive if both CPU and GPU are options
+            )
+            prompt_input = gr.Textbox(label="Positive Prompt", placeholder="e.g., a majestic lion in a vibrant jungle, photorealistic", lines=3, autofocus=True) # Autofocus on prompt
+            negative_prompt_input = gr.Textbox(label="Negative Prompt (Optional)", placeholder="e.g., blurry, low quality, deformed, watermark", lines=2)
+
+            with gr.Accordion("Advanced Settings", open=False): # Keep advanced settings initially closed
+                with gr.Row():
+                    steps_slider = gr.Slider(minimum=5, maximum=150, value=30, label="Inference Steps", step=1)
+                    cfg_slider = gr.Slider(minimum=1.0, maximum=30.0, value=7.5, label="CFG Scale", step=0.1)
+                with gr.Row():
+                     scheduler_dropdown = gr.Dropdown(
+                        choices=scheduler_choices,
+                        value=DEFAULT_SCHEDULER,
+                        label="Scheduler"
+                    )
+                     size_dropdown = gr.Dropdown(
+                        choices=SUPPORTED_SD15_SIZES,
+                        value="512x512", # SD1.5 default
+                        label="Image Size"
+                    )
+                # Combine seed input and randomize checkbox
+                with gr.Row():
+                     seed_input = gr.Number(label="Seed", value=0, precision=0, interactive=True) # Use 0 as default, interactive initially
+                     randomize_seed_checkbox = gr.Checkbox(label="Randomize seed (-1 equivalent)", value=True)
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+
+            generate_button = gr.Button("✨ Generate Image ✨", variant="primary", scale=1) # Added emojis
+
+        with gr.Column(scale=3): # Give more space to image
+            output_image = gr.Image(
+                label="Generated Image",
+                type="pil",
+                height=768, # Slightly larger preview if possible
+                width=768, # Match height for square
+                show_share_button=True,
+                show_download_button=True,
+                interactive=False # Output image is not interactive
             )
+            # The template returned the seed, let's add a display for the actual seed used
+            actual_seed_output = gr.Number(label="Actual Seed Used", precision=0, interactive=False)
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
+
+    # Link button click to generation function
+    # Use gr.on as in the template
     gr.on(
-        triggers=[run_button.click, prompt.submit],
+        triggers=[generate_button.click, prompt_input.submit], # Also trigger on prompt submit
         fn=infer,
         inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
+            model_dropdown,
+            device_dropdown,
+            prompt_input,
+            negative_prompt_input,
+            steps_slider,
+            cfg_slider,
+            scheduler_dropdown,
+            size_dropdown,
+            seed_input,
+            randomize_seed_checkbox, # Pass the checkbox value
         ],
-        outputs=[result, seed],
+        outputs=[output_image, actual_seed_output], # Return image and the actual seed used
+        api_name="generate" # Optional: For API access
     )
 
+    # Add examples from template
+    # Ensure examples match the input types and order of the infer function
+    # Examples inputs: [prompt, neg_prompt, seed, randomize_seed, width, height, cfg_scale, steps]
+    # Note: Size and Scheduler are not easily handled in standard examples, they'll use defaults.
+    # Let's adjust the example inputs to match the infer function's first few parameters
+    example_prompts = [
+        ["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", None, 0, True, "512x512", "Euler", 7.5, 30],
+        ["An astronaut riding a green horse", None, 0, True, "512x512", "Euler", 7.5, 30],
+        ["A delicious ceviche cheesecake slice", None, 0, True, "512x512", "Euler", 7.5, 30],
+    ]
+    # Update example inputs to match the infer function parameters
+    # [model_identifier, selected_device_str, prompt, negative_prompt, steps, cfg_scale, scheduler_name, size, seed, randomize_seed]
+    # Need to add dummy values for model, device, steps, cfg, scheduler, size for examples
+    # Let's simplify examples to just prompt/neg_prompt for typical template usage
+    template_examples = [
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+        "An astronaut riding a green horse",
+        "A delicious ceviche cheesecake slice",
+    ]
+    # Map template examples to the input components: [prompt_input]
+    gr.Examples(examples=template_examples, inputs=[prompt_input])
+
+
+    # Add some notes/footer from user's script, adapted for Spaces
+    gr.Markdown(
+        """
+        ---
+        **Usage Notes:**
+        1. Select a model from the dropdown (Hugging Face Hub ID). Models are downloaded and cached on the Space.
+        2. Choose your processing device (GPU recommended if available).
+        3. Enter your positive and optional negative prompts.
+        4. Adjust advanced settings (Steps, CFG Scale, Scheduler, Size, Seed) if needed. Seed -1 or the "Randomize seed" checkbox will use a random seed.
+        5. Click "Generate Image".
+        The first generation with a new model/device might take some time to load.
+        """ # Removed notes about local models and batch files
+    )
+
+
+# --- Launch the App ---
 if __name__ == "__main__":
-    demo.launch()
+    print("\n--- Starting CipherCore Stable Diffusion 1.5 Generator (Hugging Face Spaces) ---")
+    cuda_status = "CUDA available" if torch.cuda.is_available() else "CUDA not available"
+    gpu_count_str = f"Found {torch.cuda.device_count()} GPU(s)." if torch.cuda.is_available() else ""
+
+    print(f"{cuda_status} {gpu_count_str}")
+    print(f"Available devices detected by PyTorch: {', '.join(AVAILABLE_DEVICES)}")
+    print(f"Default device selected by app: {DEFAULT_DEVICE}")
+    if current_pipeline:
+         print(f"Initial model '{current_model_id}' loaded successfully.")
+    else:
+         print("No initial model loaded. Check model list and network connectivity.")
+
+
+    print("Launching Gradio interface...")
+    # For Spaces, usually launched directly without launch() parameters in app.py
+    # Spaces handles the server_name, server_port, share, etc.
+    # If running locally for testing, uncomment demo.launch()
+    # demo.launch(show_error=True, inbrowser=True) # Uncomment for local testing
+    demo.launch() # Standard launch for Hugging Face Spaces
+
+    print("Gradio interface closing.")