Update app.py

app.py

@@ -1,8 +1,6 @@
 # -*- coding: utf-8 -*-
 """
-CipherCore SD1.5 Image Generator, FAST CPU INFERENCE
-Raxephion @2025
-
+Author: @Raxephion 2025
 """
 
 import gradio as gr
@@ -46,10 +44,8 @@ DEFAULT_HUB_MODELS = [
 ]
 
 # --- Constants for Gradio UI / Generation ---
-MAX_SEED = np.iinfo(np.int32).max # <-- Added this line back! Defines the maximum seed value
-# MAX_IMAGE_SIZE = 1024 # This was in the template but not strictly used in the logic outside the slider max, keeping it defined is harmless but maybe not needed
-# However, let's add MAX_IMAGE_SIZE as it was related to the sliders in the template
-MAX_IMAGE_SIZE_SLIDER = 1024 # Renamed to avoid confusion, used only for slider max
+MAX_SEED = np.iinfo(np.int32).max # Defines the maximum seed value
+# MAX_IMAGE_SIZE_SLIDER = 1024 # Used only for slider max, kept for consistency if needed elsewhere
 
 # --- Determine available devices and set up options ---
 # This logic is from the user's script and works well for Spaces
@@ -148,28 +144,30 @@ def infer(
     cfg_scale,                # From cfg_slider
     scheduler_name,           # From scheduler_dropdown
     size,                     # From size_dropdown
-    seed,                     # From seed_input
+    seed,                     # From seed_input (now a Slider)
     randomize_seed,           # From randomize_seed_checkbox
     progress=gr.Progress(track_tqdm=True), # Added progress argument from template
 ):
     """Generates an image using the selected model and parameters on the chosen device."""
-    global current_pipeline, current_model_id, current_device_loaded, SCHEDULER_MAP, MAX_SEED # <-- Make MAX_SEED global
-
-    # Check if initial load failed or model is different
-    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):
-         # This check is done before parameter parsing so we can determine device/dtype for loading
-         # Need to redo some parameter parsing here to get device_to_use early
-         temp_device_to_use = "cuda" if selected_device_str == "GPU" and "GPU" in AVAILABLE_DEVICES else "cpu"
+    global current_pipeline, current_model_id, current_device_loaded, SCHEDULER_MAP, MAX_SEED # MAX_SEED is global
+
+    # This check is done before parameter parsing so we can determine device/dtype for loading
+    # Need to redo some parameter parsing here to get device_to_use early
+    temp_device_to_use = "cuda" if selected_device_str == "GPU" and "GPU" in AVAILABLE_DEVICES else "cpu"
+    temp_dtype_to_use = torch.float32
+    if temp_device_to_use == "cuda":
+        if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 7:
+             temp_dtype_to_use = torch.float16
+        else:
+             temp_dtype_to_use = torch.float32
+    else:
          temp_dtype_to_use = torch.float32
-         if temp_device_to_use == "cuda":
-             if torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 7:
-                  temp_dtype_to_use = torch.float16
-             else:
-                  temp_dtype_to_use = torch.float32
-         else:
-              temp_dtype_to_use = torch.float32
 
-         # Now proceed with actual model loading based on parsed device/dtype
+    # 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) != temp_device_to_use):
+
          print(f"Loading model: {model_identifier} onto {temp_device_to_use} with dtype {temp_dtype_to_use}...")
          # Clear previous pipeline to potentially free memory *before* loading the new one
          if current_pipeline is not None:
@@ -198,7 +196,7 @@ def infer(
              pipeline = StableDiffusionPipeline.from_pretrained(
                  model_identifier,
                  torch_dtype=temp_dtype_to_use, # Use the determined dtype for loading
-                 safety_checker=None,
+                 safety_checker=None, # DISABLED
              )
              pipeline = pipeline.to(temp_device_to_use) # Use the determined device
 
@@ -249,9 +247,11 @@ def infer(
     device_to_use = str(current_pipeline.device) if current_pipeline else ("cuda" if selected_device_str == "GPU" and "GPU" in AVAILABLE_DEVICES else "cpu")
     dtype_to_use = current_pipeline.dtype if current_pipeline else torch.float32 # Fallback if somehow pipeline is still None
 
-    # Check if pipeline is successfully loaded before proceeding
+
+    # Check if pipeline is successfully loaded before proceeding with generation
     if current_pipeline is None:
-         raise gr.Error("Model failed to load. Cannot generate image.")
+         raise gr.Error("Model failed to load during setup or switching. Cannot generate image.")
+
 
     # 2. Configure Scheduler
     selected_scheduler_class = SCHEDULER_MAP.get(scheduler_name)
@@ -307,38 +307,35 @@ def infer(
     # The generator device needs to match the pipeline device
     generator_device = current_pipeline.device # Must match the pipeline device
 
+    seed_int = 0 # Default if issue occurs
     if randomize_seed: # Use the randomize_seed checkbox
-        seed = random.randint(0, MAX_SEED) # Re-randomize seed using MAX_SEED
-        print(f"Randomizing seed to: {seed}")
+        seed_int = random.randint(0, MAX_SEED) # Re-randomize seed using MAX_SEED
+        print(f"Randomizing seed to: {seed_int}")
     else:
-        # Use provided seed if not randomizing (-1 will still use it)
-        # If seed is -1 but randomize_seed is False, we should still use a random seed,
-        # but make it deterministic based on system time or similar if possible, or just MAX_SEED based random.
-        # Given the UI has a 'randomize_seed' checkbox, let's treat seed=-1 as a specific "use system random" signal IF randomize_seed is also False.
-        # Or, more simply, if randomize_seed is False, always use the value in the seed input.
-        # Let's stick to the simpler interpretation: if randomize_seed is True, override seed input. If False, use seed input value.
-        print(f"Using provided seed: {int(seed)}")
+        # Use provided seed from the slider input
+        try:
+             seed_int = int(seed)
+             print(f"Using provided seed: {seed_int}")
+        except ValueError:
+             print(f"Warning: Invalid seed input '{seed}'. Using random seed instead.")
+             gr.Warning(f"Invalid seed input '{seed}'. Using random seed instead.")
+             seed_int = random.randint(0, MAX_SEED) # Fallback to random if input is not int
+             randomize_seed = True # Mark as randomized for reporting
 
 
+    generator = None # Initialize generator as None
     try:
-        # Ensure seed is an integer for the generator
-        seed_int = int(seed)
         # Explicitly move generator to the desired device
         generator = torch.Generator(device=generator_device).manual_seed(seed_int)
         print(f"Generator set with seed {seed_int} on device: {generator_device}")
     except Exception as e:
-         # Handle potential issues like non-integer seed input
-         print(f"Warning: Error setting seed generator with seed {seed} on device {generator_device}: {e}. Falling back to default generator (potentially on CPU) and using a potentially different seed.")
-         gr.Warning(f"Failed to set seed generator with seed {seed}. Using random seed. Error: {e}")
+         print(f"Warning: Error setting seed generator with seed {seed_int} on device {generator_device}: {e}. Falling back to default generator (potentially on CPU) or system random.")
+         gr.Warning(f"Failed to set seed generator with seed {seed_int}. Using system random seed. Error: {e}")
          generator = None # Let pipeline handle random seed if generator creation fails or device mismatch
-         # If generator creation failed, the actual seed used might be different. How to report it?
-         # The simplest is to let the pipeline run and report the seed *we attempted to use*.
-         # Or, if we know pipeline will use a random seed on failure, report that.
-         # For simplicity, let's report the seed_int we derived UNLESS randomize_seed was True, then report the generated random seed.
-         # No, the most accurate is to let the pipeline generate, and some diffusers versions return the generator state/seed.
-         # However, the standard pipeline output doesn't return the seed.
-         # Let's just report the seed we *tried* to use. If randomize_seed was True, it's the random one. If False, it's the user input one.
-         pass # Keep the last calculated seed_int or seed value
+         # If generator creation failed, the actual seed used by the pipeline will be different and system-dependent random.
+         # We should probably report -1 in this case, or just report the seed we tried to use.
+         # Reporting the seed we *tried* to use is simpler and often sufficient.
+         pass # Keep the last calculated seed_int
 
     # 5. Generate Image
     # Ensure required parameters are integers/floats
@@ -351,7 +348,7 @@ def infer(
     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={seed_int if generator else 'Random (Generator Failed)'}, Device={device_to_use}, Dtype={dtype_to_use}")
+    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={seed_int if generator else 'System Random'}, Device={device_to_use}, Dtype={dtype_to_use}")
     start_time = time.time()
 
     try:
@@ -378,8 +375,8 @@ def infer(
         print(f"Generation finished in {end_time - start_time:.2f} seconds.")
         generated_image = output.images[0]
 
-        # Determine the seed to return: the one we attempted to use
-        actual_seed_used = seed_int if generator else -1 # Return -1 or the input seed if generator failed
+        # Determine the seed to return: the one we attempted to use, or -1 if generator creation failed
+        actual_seed_used = seed_int if generator else -1
 
         # Return both the image and the seed (potentially randomized)
         return generated_image, actual_seed_used
@@ -422,17 +419,9 @@ if not model_choices:
 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 # Make it interactive if there's *any* model choice
-
-# Ensure the initial default model is actually in the choices list if possible
-if initial_default_model != "No models found" and initial_default_model not in initial_model_choices:
-     print(f"Warning: Initial default model '{initial_default_model}' is not in the model_choices list.")
-     if initial_model_choices and initial_model_choices[0] != "No models found":
-          initial_default_model = initial_model_choices[0]
-          print(f"Setting default model to first available choice: {initial_default_model}")
-     else:
-          initial_default_model = "No models found" # Fallback if no choices
+    initial_default_model = INITIAL_MODEL_ID if INITIAL_MODEL_ID and INITIAL_MODEL_ID in initial_model_choices else (initial_model_choices[0] if initial_model_choices and initial_model_choices[0] != "No models found" else "No models found")
+
+    model_dropdown_interactive = True if initial_default_model != "No models found" else False # Make it interactive if there's *any* model choice
 
 scheduler_choices = list(SCHEDULER_MAP.keys())
 
@@ -455,10 +444,8 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: # Added Soft theme from
     )
 
     # 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.)*")
-    elif initial_default_model != "No models found":
-         gr.Markdown(f"*(Note: Loading model '{initial_default_model}' on first generation... This might take some time.)*")
+    if initial_default_model != "No models found":
+         gr.Markdown(f"*(Note: Model '{initial_default_model}' will load on first generation or model switch... This might take some time.)*")
     else:
          gr.Markdown(f"*(Note: No models available. Add Hub IDs to DEFAULT_HUB_MODELS in the script.)*")
 
@@ -497,8 +484,8 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo: # Added Soft theme from
                     )
                 # Combine seed input and randomize checkbox
                 with gr.Row():
-                     # Use MAX_SEED for slider max
-                     seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, precision=0, interactive=True) # Use 0 as default, interactive initially
+                     # Removed precision=0 from Slider
+                     seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, interactive=True) # Use 0 as default, interactive initially
                      randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True) # Simplified label
 
 
@@ -578,7 +565,7 @@ if __name__ == "__main__":
     if current_pipeline:
          print(f"Initial model '{current_model_id}' loaded successfully on {current_device_loaded}.")
     else:
-         print("No initial model loaded. Check model list and network connectivity.")
+         print("No initial model loaded or initial load failed. Check model list and network connectivity.")
 
 
     print("Launching Gradio interface...")