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 # 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_seg = StableDiffusionControlNetPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", controlnet=controlnet_seg, torch_dtype=torch.float16).to("cuda") pipe_depth = StableDiffusionControlNetPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", controlnet=controlnet_depth, torch_dtype=torch.float16).to("cuda") @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") 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="", text_input="the main object"): """ Prompts Florence-2 to isolate an object and map out a clean segmentation mask. """ if task_prompt == "": 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") #image_np = np.array(init_image) conditioning_image = get_florence_segmentation( init_image, task_prompt="", 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") cv_img = np.array(pil_image) 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] # ["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) 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] 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)