app.py

import gradio as gr
import spaces
import torch
import os
import gradio as gr
import cv2
import numpy as np
from diffusers.utils import load_image 
from PIL import Image
import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler
import gc
import torch

# Programmatically add the repository to PyTorch Hub's trusted list
#torch.hub.set_dir('/root/.cache/torch/hub') # Optional: ensures clean cache path
#torch.hub.trusted_list.append("rwightman/gen-efficientnet-pytorch")
#torch.hub.trusted_list.append("intel-isl/MiDaS")
import torch.hub

# Globally override the default 'check' state to always force allow downloads
orig_load = torch.hub.load
def patched_load(*args, **kwargs):
    kwargs['trust_repo'] = True
    return orig_load(*args, **kwargs)

torch.hub.load = patched_load


device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

print("Initializing Models (Florence-2 & )...")
florence_model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True, torch_dtype=torch_dtype).to(device).eval()
florence_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)

# Load a quick, highly accurate MiDaS model pipeline via torch hub
midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small",trust_repo=True).to(device).eval()
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms",trust_repo=True)
transform_depth = midas_transforms.flag_transform if hasattr(midas_transforms, 'flag_transform') else midas_transforms.dpt_transform
    

# 1. Load the specific ControlNet adapter from Hugging Face
controlnet_ids= ["lllyasviel/sd-controlnet-canny","lllyasviel/sd-controlnet-seg","lllyasviel/sd-controlnet-depth"]

controlnet_edge = ControlNetModel.from_pretrained(controlnet_ids[0],torch_dtype=torch.float16)
controlnet_seg = ControlNetModel.from_pretrained(controlnet_ids[1],torch_dtype=torch.float16)
controlnet_depth = ControlNetModel.from_pretrained(controlnet_ids[2],torch_dtype=torch.float16)


pipe_edge = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet_edge, torch_dtype=torch.float16).to("cuda")

pipe_edge.enable_attention_slicing()

pipe_seg = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet_seg, torch_dtype=torch.float16).to("cuda")
pipe_seg.enable_attention_slicing()


pipe_depth = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet_depth, torch_dtype=torch.float16).to("cuda")
pipe_depth.enable_attention_slicing()




@spaces.GPU
def greet(n):
    print(zero.device) # <-- 'cuda:0' 🤗
    return f"Hello {zero + n} Tensor"


IMAGE_FOLDER = "./images"  # Change this to your folder path

@spaces.GPU
def generate_edgeImage(image_path):
    # Convert image to numpy array to generate Canny edges
    init_image = Image.open(image_path).convert("RGB")
    init_image.thumbnail((512, 512), Image.Resampling.LANCZOS)
    image_np = np.array(init_image)
    low_threshold = 100
    high_threshold = 200
    edges = cv2.Canny(image_np, low_threshold, high_threshold)
    # Convert back to PIL format for the pipeline
    canny_image = Image.fromarray(edges).convert("RGB")
    return canny_image
    
@spaces.GPU
def get_florence_segmentation(image, task_prompt="<REFERRING_EXPRESSION_SEGMENTATION>", text_input="the main object"):
    """
    Prompts Florence-2 to isolate an object and map out a clean segmentation mask.
    """
    if task_prompt == "<REFERRING_EXPRESSION_SEGMENTATION>":
        prompt = task_prompt + text_input
    else:
        prompt = task_prompt

    inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)

    with torch.no_grad():
        generated_ids = florence_model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=1024,
            num_beams=3
        )

    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
    parsed_answer = florence_processor.post_process_generation(
        generated_text,
        task=task_prompt,
        image_size=(image.width, image.height)
    )

    # Extract polygon arrays and build a solid binary conditioning mask image
    mask = Image.new("RGB", image.size, "black")
    try:
        # Pull polygons returned from the segmentation prompt
        polygons = parsed_answer[task_prompt]['polygons']
        from PIL import ImageDraw
        draw = ImageDraw.Draw(mask)
        for poly in polygons:
            # Flatten array for PIL polygon layout [x1, y1, x2, y2...]
            flat_poly = [coord for point in poly for coord in point]
            if len(flat_poly) >= 6:
                draw.polygon(flat_poly, fill="white")
    except KeyError:
        print("Target object segment not clearly resolved. Using fallback empty mask.")

    return mask


@spaces.GPU
def generate_segmt_img(image_path,generate_segmt_img):
    init_image = Image.open(image_path).convert("RGB")
    init_image.thumbnail((512, 512), Image.Resampling.LANCZOS)
    #image_np = np.array(init_image)
    
    conditioning_image = get_florence_segmentation(
        init_image, task_prompt="<REFERRING_EXPRESSION_SEGMENTATION>",
        text_input=generate_segmt_img)
    return conditioning_image

@spaces.GPU
def extract_depth_map(image_path):
    """
    Utilizes a lightweight MiDaS pipeline to extract depth estimation layouts.
    """
    pil_image = Image.open(image_path).convert("RGB")
    pil_image.thumbnail((512, 512), Image.Resampling.LANCZOS)
    cv_img = np.array(pil_image)
    print(f"cv_img: '{cv_img.shape}'")
    img_tensor = transform_depth(cv_img).to(device)

    with torch.no_grad():
        prediction = midas(img_tensor)
        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=pil_image.size[::-1],
            mode="bicubic",
            align_corners=False,
        ).squeeze()

    depth_output = prediction.cpu().numpy()
    # Normalize pixel depth values into a visible 0-255 map array
    depth_min, depth_max = depth_output.min(), depth_output.max()
    normalized_depth = (255 * (depth_output - depth_min) / (depth_max - depth_min)).astype(np.uint8)
    depth_rgb = np.concatenate([normalized_depth[:, :, None]] * 3, axis=2)
    return Image.fromarray(depth_rgb)
    
@spaces.GPU   
def get_image_list(folder):
    if not os.path.exists(folder):
        return []
    valid_extensions = (".png", ".jpg", ".jpeg", ".webp", ".gif")
    return [f for f in os.listdir(folder) if f.lower().endswith(valid_extensions)]

image_files = get_image_list(IMAGE_FOLDER)

@spaces.GPU
def load_image(selected_file):
    if not selected_file:
        return None
    return os.path.join(IMAGE_FOLDER, selected_file)

# generate input image 
@spaces.GPU
def generate_input_image(image_path,input_choice,seg_input_promt):
    print(f"image_path: '{image_path}'")
    print(f"input_choice: '{input_choice}'")
    image_path = image_path[0]
    #gc.collect()
    #torch.cuda.empty_cache()
    # ["Edge image","Segmented image","Depth Image"]
    if input_choice == "Edge image" :
        conditioning_image = generate_edgeImage(image_path)
    elif input_choice == "Segmented image" :
        conditioning_image = generate_segmt_img(image_path,seg_input_promt)
    else:
        conditioning_image = extract_depth_map(image_path)
    gc.collect()
    torch.cuda.empty_cache()
    return conditioning_image

@spaces.GPU
def generate_output_image(prompt,conditioning_image,input_choice):
    #---------------------
    if input_choice == "Edge image":
        pipe = pipe_edge
    elif input_choice == "Segmented image" :
        pipe = pipe_seg
    else:
        pipe = pipe_depth
        
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
    pipe.enable_attention_slicing()
    #---------------------
    #prompt = "A majestic fantasy castle built into a geometric glass mountain, sharp focus, hyper-detailed digital art"
    negative_prompt = "blurry, low quality, human, text, worst composition"

    generated_output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=conditioning_image,
        num_inference_steps=20,
        guidance_scale=8.0
        ).images[0]
    gc.collect()
    torch.cuda.empty_cache()
    return generated_output
        
    
# Generation pipeline function using both the source image and text prompt
@spaces.GPU
def process_and_generate(selected_file, prompt):
    if not selected_file:
        return None, "Please select an input image first."
    if not prompt:
        return None, "Please enter a text prompt."
    
    input_path = os.path.join(IMAGE_FOLDER, selected_file)
    
    # -------------------------------------------------------------
    # PLACEHOLDER: Insert your Multi-ControlNet / SD pipeline here
    # e.g., output = pipeline(image=input_path, prompt=prompt)
    # -------------------------------------------------------------
    
    # For demonstration, returning the input image path
    generated_output = input_path 
    
    return generated_output, f"Generated matching prompt: '{prompt}'"

with gr.Blocks() as demo:
    gr.Markdown("# Guided Image Generation Pipeline")
    
    with gr.Row():
        # Column 1: Source Image Selection
        with gr.Column():
            gr.Markdown("### 1. Source Image")
            file_input = gr.File(
                label="Select Images or a Folder", 
                file_count="multiple", 
                file_types=["image"]
            )
            seg_input = gr.Textbox( label="segmentation Prompt", 
                placeholder="Describe modifications, styles, or Canny/Depth map directions...")
            with gr.Row():
                image_dropdown = gr.Dropdown(
                    choices=["Edge image","Segmented image","Depth Image"] , 
                    label="Select Input image method",
                    value="Edge image"
                )
                generate_inp_btn = gr.Button("Generate Input", variant="primary")
            
            input_preview = gr.Image(type="pil", label="Input Preview")
            
        # Column 2: Prompt Engineering & Execution
        with gr.Column():
            gr.Markdown("### 2. Generation Settings")
            prompt_input = gr.Textbox(
                label="Text Prompt", 
                placeholder="Describe modifications, styles, or Canny/Depth map directions...",
                lines=3
            )
            generate_btn = gr.Button("Generate Image", variant="primary")
            
        # Column 3: Output Visualization
        with gr.Column():
            gr.Markdown("### 3. Pipeline Output")
            output_image = gr.Image(type="filepath", label="Generated Output")
            status_text = gr.Textbox(label="Status", interactive=False)
            
    # Core event binds
    #image_dropdown.change(
    #    fn=load_image, 
    #    inputs=image_dropdown, 
    #    outputs=input_preview
    #)
    generate_inp_btn.click(
        fn=generate_input_image,
        inputs=[file_input,image_dropdown,seg_input],
        outputs=[input_preview]
    )
    generate_btn.click(
        fn=generate_output_image,
        inputs=[prompt_input,input_preview,image_dropdown],
        outputs=[output_image]
    )
    
    # Initial state setup
    #if image_files:
    #    demo.load(
    #        fn=load_image, 
    #        inputs=image_dropdown, 
    #        outputs=input_preview
    #    )

if __name__ == "__main__":
    demo.launch(share=True)