Update app.py

app.py

@@ -1,118 +1,49 @@
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import DiffusionPipeline
 import gradio as gr
-import spaces
-import os
-import sys
-import subprocess
-import numpy as np
 from PIL import Image
-import cv2
-
-import torch
-
-from diffusers import StableDiffusion3ControlNetPipeline
-from diffusers.models import SD3ControlNetModel, SD3MultiControlNetModel
-from diffusers.utils import load_image
-
-# load pipeline
-controlnet_canny = SD3ControlNetModel.from_pretrained("InstantX/SD3-Controlnet-Canny")
-pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-3-medium-diffusers",
-    controlnet=controlnet_canny
-).to("cuda", torch.float16)
-
-def resize_image(input_path, output_path, target_height):
-    # Open the input image
-    img = Image.open(input_path)
-
-    # Calculate the aspect ratio of the original image
-    original_width, original_height = img.size
-    original_aspect_ratio = original_width / original_height
-
-    # Calculate the new width while maintaining the aspect ratio and the target height
-    new_width = int(target_height * original_aspect_ratio)
 
-    # Resize the image while maintaining the aspect ratio and fixing the height
-    img = img.resize((new_width, target_height), Image.LANCZOS)
+# Load Stable Diffusion 3 (from InstantX)
+model_id = "instantx/stable-diffusion-3-medium"
 
-    # Save the resized image
-    img.save(output_path)
+# Load the ControlNet model (use an appropriate pre-trained controlnet model)
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
 
-    return output_path, new_width, target_height
+# Set up the pipeline using both SD3 and ControlNet
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    model_id,
+    controlnet=controlnet,
+    torch_dtype=torch.float16
+)
 
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe.to(device)
 
-@spaces.GPU(duration=90)
-def infer(image_in, prompt, inference_steps, guidance_scale, control_weight, progress=gr.Progress(track_tqdm=True)):
-
-    n_prompt = 'NSFW, nude, naked, porn, ugly'
-
-    # Canny preprocessing
-    image_to_canny = load_image(image_in)
-    image_to_canny = np.array(image_to_canny)
-    image_to_canny = cv2.Canny(image_to_canny, 100, 200)
-    image_to_canny = image_to_canny[:, :, None]
-    image_to_canny = np.concatenate([image_to_canny, image_to_canny, image_to_canny], axis=2)
-    image_to_canny = Image.fromarray(image_to_canny)
-
-    control_image = image_to_canny
- 
-    # infer
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=n_prompt,
-        control_image=control_image, 
-        controlnet_conditioning_scale=control_weight,
-        num_inference_steps=inference_steps,
-        guidance_scale=guidance_scale,
-    ).images[0]
-
+# Function for Img2Img with ControlNet
+def controlnet_img2img(image, prompt, strength=0.8, guidance=7.5):
+    image = Image.fromarray(image).convert("RGB")  # Convert to RGB
     
-    image_redim, w, h = resize_image(image_in, "resized_input.jpg", 1024)
-    image = image.resize((w, h), Image.LANCZOS)
+    # Run the pipeline
+    result = pipe(prompt=prompt, image=image, strength=strength, guidance_scale=guidance).images[0]
+    return result
+
+# Gradio Interface
+def img_editor(input_image, prompt):
+    result = controlnet_img2img(input_image, prompt)
+    return result
+
+# Create Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Img2Img Editor with ControlNet and Stable Diffusion 3")
+    with gr.Row():
+        image_input = gr.Image(source="upload", type="numpy", label="Input Image")
+        prompt_input = gr.Textbox(label="Prompt")
+    result_output = gr.Image(label="Output Image")
     
-    return image, gr.update(value=image_to_canny, visible=True)
-   
-
-css="""
-#col-container{
-    margin: 0 auto;
-    max-width: 1080px;
-}
-"""
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("""
-        # SD3 ControlNet
-
-        Experiment with Stable Diffusion 3 ControlNet models proposed and maintained by the InstantX team.<br />
-        Model card: [InstantX/SD3-Controlnet-Canny](https://huggingface.co/InstantX/SD3-Controlnet-Canny)
-        """)
-        
-        with gr.Column():
-            
-            with gr.Row():
-                with gr.Column():
-                    image_in = gr.Image(label="Image reference", sources=["upload"], type="filepath")
-                    prompt = gr.Textbox(label="Prompt")
-                    
-                    with gr.Accordion("Advanced settings", open=False):
-                        with gr.Column():
-                            with gr.Row():
-                                inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=25)
-                                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=7.0)
-                            control_weight = gr.Slider(label="Control Weight", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
-                    
-                    submit_canny_btn = gr.Button("Submit")
-                    
-                with gr.Column():
-                    result = gr.Image(label="Result")
-                    canny_used = gr.Image(label="Preprocessed Canny", visible=False)
-                    
-
+    submit_btn = gr.Button("Generate")
+    submit_btn.click(fn=img_editor, inputs=[image_input, prompt_input], outputs=result_output)
 
-    submit_canny_btn.click(
-        fn = infer,
-        inputs = [image_in, prompt, inference_steps, guidance_scale, control_weight],
-        outputs = [result, canny_used],
-        show_api=False
-    )
-demo.queue().launch()
+# Launch Gradio interface
+demo.launch()