Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import torch
 import numpy as np
-import cv2
 from diffusers import StableDiffusionPipeline, UniPCMultistepScheduler
 from model import UNet2DConditionModelEx
 from pipeline import StableDiffusionControlLoraV3Pipeline 
@@ -11,189 +10,229 @@ from huggingface_hub import login
 import spaces
 import random
 from pathlib import Path
+import hashlib
+import datetime
+import json
+from tqdm import tqdm
 
 # 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"
-
-# Load the custom UNet
-unet = UNet2DConditionModelEx.from_pretrained(
-    base_model,
-    subfolder="unet"
-)
-
-unet = unet.add_extra_conditions("ow-gbi-control-lora")
-
-pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
-    base_model, 
-    unet=unet
-)
-
-# Performance optimizations
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-pipe.enable_attention_slicing()
-pipe.enable_vae_slicing()
-
-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]"))
-        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)
-    elif isinstance(image, str):
-        image = np.array(Image.open(image))
-    
-    if len(image.shape) == 2:
-        image = np.stack([image] * 3, axis=-1)
-    elif image.shape[2] == 4:
-        image = image[..., :3]
-    
-    canny_image = cv2.Canny(image, low_threshold, high_threshold)
-    canny_image = np.stack([canny_image] * 3, axis=-1)
-    return Image.fromarray(canny_image)
+# Setup directories
+HF_SPACE_ID = "naonauno/groundbi-factory"
+OUTPUT_DIR = "/home/user/outputs"
+
+os.makedirs('outputs', exist_ok=True)
+os.makedirs('metadata', exist_ok=True)
+metadata_dir = 'metadata'
+
+class AdvancedGenerationTracker:
+    def __init__(self, total_steps):
+        self.progress_bar = tqdm(total=total_steps, desc="Image Generation")
+        self.current_step = 0
+        self.memory_usage_log = []
+
+    def update_progress(self, step_size=1):
+        self.current_step += step_size
+        self.progress_bar.update(step_size)
+        self._log_memory_usage()
+
+    def _log_memory_usage(self):
+        if torch.cuda.is_available():
+            memory_info = {
+                'step': self.current_step,
+                'cuda_allocated': torch.cuda.memory_allocated(),
+                'cuda_reserved': torch.cuda.memory_reserved(),
+                'cuda_max_allocated': torch.cuda.max_memory_allocated()
+            }
+            self.memory_usage_log.append(memory_info)
+
+    def finalize(self):
+        self.progress_bar.close()
+        return self.memory_usage_log
+
+def setup_pipeline():
+    unet = UNet2DConditionModelEx.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        subfolder="unet"
+    )
+    unet = unet.add_extra_conditions("ow-gbi-control-lora")
+
+    pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", 
+        unet=unet
+    )
+
+    # Performance optimizations
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_attention_slicing()
+    pipe.enable_vae_slicing()
 
-@spaces.GPU(duration=180)  # Reduced to 3 minutes
-def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold, seed, progress=gr.Progress()):
-    if input_image is None:
+    pipe.load_lora_weights(
+        "models",
+        weight_name="40kHalf.safetensors"
+    )
+    return pipe
+
+pipe = setup_pipeline()
+
+def save_to_space(image, filename):
+    path = os.path.join(OUTPUT_DIR, filename)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    image.save(path)
+    return path
+
+def generate_advanced_filename(prompt, seed, style=None):
+    hash_input = f"{prompt}_{seed}"
+    filename_hash = hashlib.md5(hash_input.encode()).hexdigest()[:8]
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    style_prefix = f"{style}_" if style else ""
+    return f"{style_prefix}{timestamp}_{filename_hash}"
+
+def export_generation_metadata(metadata, output_path):
+    with open(output_path, 'w') as f:
+        json.dump(metadata, f, indent=2)
+    return output_path
+
+@spaces.GPU(duration=180)
+def generate_image(
+    image,
+    prompt,
+    negative_prompt,
+    guidance_scale,
+    steps,
+    seed,
+    strength=0.8,
+    num_images=1,
+    progress=gr.Progress()
+):
+    if image is None:
         raise gr.Error("Please provide an input image!")
     
     try:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_base_dir = os.path.join('outputs', timestamp)
+        os.makedirs(output_base_dir, exist_ok=True)
+
         if seed is not None and seed != "":
             try:
                 generator = torch.Generator().manual_seed(int(seed))
+                current_seed = int(seed)
             except ValueError:
                 generator = torch.Generator()
+                current_seed = random.randint(1, 1000000)
         else:
             generator = torch.Generator()
+            current_seed = random.randint(1, 1000000)
+
+        tracker = AdvancedGenerationTracker(steps)
+
+        def callback_on_step_end(pipeline, step, timestep, callback_kwargs):
+            tracker.update_progress()
+            if progress is not None:
+                progress(step/steps)
+            return {}
 
-        progress(0.1, desc="Processing input image...")
-        canny_image = get_canny_image(input_image, low_threshold, high_threshold)
-        
         progress(0.3, desc="Generating image...")
         with torch.no_grad():
-            image = pipe(
+            result = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 num_inference_steps=int(steps),
                 guidance_scale=float(guidance_scale),
-                image=canny_image,
+                image=image,
+                strength=strength,
                 extra_condition_scale=1.0,
-                generator=generator
-            ).images[0]
+                generator=generator,
+                num_images_per_prompt=num_images,
+                callback_on_step_end=callback_on_step_end
+            )
+
+        generated_image = result.images[0]
         
+        # Save the image
+        filename = generate_advanced_filename(prompt, current_seed)
+        image_path = os.path.join(output_base_dir, f"{filename}.png")
+        generated_image.save(image_path)
+        save_to_space(generated_image, f"{filename}.png")
+
+        # Save metadata
+        generation_metadata = {
+            "generation_timestamp": timestamp,
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "seed": current_seed,
+            "generation_parameters": {
+                "guidance_scale": guidance_scale,
+                "steps": steps,
+                "strength": strength,
+                "num_images": num_images
+            },
+            "image_file": os.path.basename(image_path)
+        }
+
+        metadata_path = os.path.join(metadata_dir, f"{filename}_metadata.json")
+        export_generation_metadata(generation_metadata, metadata_path)
+
+        memory_log = tracker.finalize()
         progress(1.0, desc="Done!")
-        return canny_image, image
-    
+        
+        return generated_image
+
     except Exception as e:
         raise gr.Error(f"An error occurred: {str(e)}")
 
-def random_image_click():
-    image_path = get_random_condition_image()
-    if image_path:
-        return Image.open(image_path)
-    return None
-
-# Example data with reduced steps
-examples = [
-    [
-        "conditions/example1.jpg",
-        "a futuristic cyberpunk city",
-        "blurry, bad quality",
-        7.5,
-        25,  # Reduced steps
-        100,
-        200,
-        42
-    ],
-    [
-        "conditions/example2.jpg",
-        "a serene mountain landscape",
-        "dark, gloomy",
-        7.0,
-        25,  # Reduced steps
-        120,
-        180,
-        123
-    ]
-]
+css = """
+.container { max-width: 900px; margin: auto; }
+.parameter-hint { font-size: 0.8em; color: #666; margin-top: -5px; }
+"""
 
 # Create the Gradio interface
-with gr.Blocks() as demo:
+with gr.Blocks(css=css) as demo:
     gr.Markdown(
         """
-        # Control LoRA v3 Demo
-        ⚠️ Warning: This is a demo of Control LoRA v3. Generation might take a few minutes.
-        For better results with ZeroGPU, it's recommended to use 20-30 steps.
-        The model uses edge detection to guide the image generation process.
+        # Terrain Generator
+        ⚠️ Warning: This is a demo running on ZeroGPU. Generation might take a few minutes.
+        For best results, use 15-20 steps for generation.
         """
     )
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Input Image", type="numpy")
-            random_image_btn = gr.Button("Load Random Reference Image")
-            status_text = gr.Textbox(label="Status", value="Ready", interactive=False)
+            input_image = gr.Image(label="Input Image", type="pil")
             
             prompt = gr.Textbox(
                 label="Prompt",
-                placeholder="Enter your prompt here... (e.g., 'a futuristic cyberpunk city')"
+                placeholder="Describe the terrain..."
             )
             negative_prompt = gr.Textbox(
                 label="Negative Prompt",
-                placeholder="Enter things you don't want to see... (e.g., 'blurry, bad quality')"
+                placeholder="What to avoid..."
+            )
+            guidance_scale = gr.Slider(
+                label="Guidance Scale", 
+                minimum=1, 
+                maximum=20, 
+                value=7.5, 
+                info="Higher = more prompt adherence, Lower = more creativity"
+            )
+            steps = gr.Slider(
+                label="Steps", 
+                minimum=1, 
+                maximum=50, 
+                value=20, 
+                info="More steps = higher quality but slower"
+            )
+            seed = gr.Textbox(
+                label="Seed (empty for random)", 
+                placeholder="Enter a number for reproducible results",
+                info="Controls randomness. Same seed = same output."
             )
-            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=50, value=25, label="Steps")  # Reduced max 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
-    )
-
-    random_image_btn.click(
-        fn=random_image_click,
-        outputs=input_image
-    )
-
     generate.click(
         fn=generate_image,
         inputs=[
@@ -202,11 +241,9 @@ with gr.Blocks() as demo:
             negative_prompt,
             guidance_scale,
             steps,
-            low_threshold,
-            high_threshold,
             seed
         ],
-        outputs=[canny_output, result]
+        outputs=result
     )
 
 demo.queue()