app.py

import gradio as gr
import torch
import numpy as np
from diffusers import StableDiffusionPipeline, UniPCMultistepScheduler
from model import UNet2DConditionModelEx
from pipeline import StableDiffusionControlLoraV3Pipeline 
from PIL import Image
import os
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"))

# 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()

    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.3, desc="Generating image...")
        with torch.no_grad():
            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=int(steps),
                guidance_scale=float(guidance_scale),
                image=image,
                strength=strength,
                extra_condition_scale=1.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 generated_image

    except Exception as e:
        raise gr.Error(f"An error occurred: {str(e)}")

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(css=css) as demo:
    gr.Markdown(
        """
        # 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="pil")
            
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the terrain..."
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                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."
            )
            generate = gr.Button("Generate")
        
        with gr.Column():
            result = gr.Image(label="Generated Image")
    
    generate.click(
        fn=generate_image,
        inputs=[
            input_image,
            prompt,
            negative_prompt,
            guidance_scale,
            steps,
            seed
        ],
        outputs=result
    )

demo.queue()
demo.launch(share=True)