Update app.py

app.py

@@ -9,13 +9,19 @@ from PIL import Image
 import os
 from huggingface_hub import login
 import spaces
+import random
+from pathlib import Path
 
 # Login using the token
 login(token=os.environ.get("HF_TOKEN"))
 
+# For deterministic generation
+torch.manual_seed(42)
+torch.backends.cudnn.deterministic = True
+
 # Initialize the models
 base_model = "runwayml/stable-diffusion-v1-5"
-dtype = torch.float16  # A100 works better with float16
+dtype = torch.float16
 
 # Load the custom UNet
 unet = UNet2DConditionModelEx.from_pretrained(
@@ -24,25 +30,30 @@ unet = UNet2DConditionModelEx.from_pretrained(
     torch_dtype=dtype
 )
 
-# Add conditioning with ow-gbi-control-lora
 unet = unet.add_extra_conditions("ow-gbi-control-lora")
 
-# Create the pipeline with custom UNet
 pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
     base_model, 
     unet=unet,
     torch_dtype=dtype
 )
 
-# Use a faster scheduler
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
-# Load the ControlLoRA weights
 pipe.load_lora_weights(
     "models",
     weight_name="40kHalf.safetensors"
 )
 
+def get_random_condition_image():
+    conditions_dir = Path("conditions")
+    if conditions_dir.exists():
+        image_files = list(conditions_dir.glob("*.[jp][pn][g]"))  # matches .jpg, .png, .jpeg
+        if image_files:
+            random_image = random.choice(image_files)
+            return str(random_image)
+    return None
+
 def get_canny_image(image, low_threshold=100, high_threshold=200):
     if isinstance(image, Image.Image):
         image = np.array(image)
@@ -54,8 +65,16 @@ def get_canny_image(image, low_threshold=100, high_threshold=200):
     canny_image = np.stack([canny_image] * 3, axis=-1)
     return Image.fromarray(canny_image)
 
-@spaces.GPU(duration=120)  # Set GPU allocation duration to 120 seconds
-def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold):
+@spaces.GPU(duration=120)
+def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold, seed):
+    if seed is not None and seed != "":
+        try:
+            generator = torch.Generator().manual_seed(int(seed))
+        except ValueError:
+            generator = torch.Generator()
+    else:
+        generator = torch.Generator()
+
     canny_image = get_canny_image(input_image, low_threshold, high_threshold)
     
     with torch.no_grad():
@@ -65,29 +84,102 @@ def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps,
             num_inference_steps=steps,
             guidance_scale=guidance_scale,
             image=canny_image,
-            extra_condition_scale=1.0
+            extra_condition_scale=1.0,
+            generator=generator
         ).images[0]
     
     return canny_image, image
 
+def random_image_click():
+    image_path = get_random_condition_image()
+    if image_path:
+        return Image.open(image_path)
+    return None
+
+# Example data
+examples = [
+    [
+        "conditions/example1.jpg",  # Replace with actual paths
+        "a futuristic cyberpunk city",
+        "blurry, bad quality",
+        7.5,
+        50,
+        100,
+        200,
+        42
+    ],
+    [
+        "conditions/example2.jpg",  # Replace with actual paths
+        "a serene mountain landscape",
+        "dark, gloomy",
+        7.0,
+        40,
+        120,
+        180,
+        123
+    ]
+]
+
 # Create the Gradio interface
 with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # Control LoRA v3 Demo
+        ⚠️ Warning: This is a demo of Control LoRA v3. Please be aware that generation can take several minutes. 
+        The model uses edge detection to guide the image generation process.
+        """
+    )
+    
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="numpy")
-            prompt = gr.Textbox(label="Prompt")
-            negative_prompt = gr.Textbox(label="Negative Prompt")
+            random_image_btn = gr.Button("Load Random Reference Image")
+            
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your prompt here... (e.g., 'a futuristic cyberpunk city')"
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="Enter things you don't want to see... (e.g., 'blurry, bad quality')"
+            )
             with gr.Row():
                 low_threshold = gr.Slider(minimum=1, maximum=255, value=100, label="Canny Low Threshold")
                 high_threshold = gr.Slider(minimum=1, maximum=255, value=200, label="Canny High Threshold")
             guidance_scale = gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale")
             steps = gr.Slider(minimum=1, maximum=100, value=50, label="Steps")
+            seed = gr.Textbox(label="Seed (empty for random)", placeholder="Enter a number for reproducible results")
             generate = gr.Button("Generate")
         
         with gr.Column():
             canny_output = gr.Image(label="Canny Edge Detection")
             result = gr.Image(label="Generated Image")
     
+    # Set up example gallery
+    gr.Examples(
+        examples=examples,
+        inputs=[
+            input_image,
+            prompt,
+            negative_prompt,
+            guidance_scale,
+            steps,
+            low_threshold,
+            high_threshold,
+            seed
+        ],
+        outputs=[canny_output, result],
+        fn=generate_image,
+        cache_examples=True
+    )
+
+    # Handle the random image button
+    random_image_btn.click(
+        fn=random_image_click,
+        outputs=input_image
+    )
+
+    # Handle the generate button
     generate.click(
         fn=generate_image,
         inputs=[
@@ -97,7 +189,8 @@ with gr.Blocks() as demo:
             guidance_scale,
             steps,
             low_threshold,
-            high_threshold
+            high_threshold,
+            seed
         ],
         outputs=[canny_output, result]
     )