Update app.py

app.py

@@ -1,50 +1,78 @@
-import gradio as gr
-import torch
-from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline
 import spaces
+from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, EulerAncestralDiscreteScheduler
+import torch
+import gradio as gr
 from PIL import Image
 import numpy as np
 
-# Load the ControlNet model and pipeline
+# Load the models
 controlnet = ControlNetModel.from_pretrained(
     "briaai/BRIA-2.2-ControlNet-Recoloring",
     torch_dtype=torch.float16
-)
+).to('cuda')
+
 pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     "briaai/BRIA-2.2",
     controlnet=controlnet,
     torch_dtype=torch.float16,
-).to("cuda")
+    device_map='auto',
+    low_cpu_mem_usage=True,
+    offload_state_dict=True,
+).to('cuda')
+
+pipe.scheduler = EulerAncestralDiscreteScheduler(
+    beta_start=0.00085,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    num_train_timesteps=1000,
+    steps_offset=1
+)
+
+pipe.force_zeros_for_empty_prompt = False
+
+def resize_image(image):
+    image = image.convert('RGB')
+    current_size = image.size
+    transform = gr.Image(height=1024, width=1024, keep_aspect_ratio=True, source="upload", tool="editor")
+    resized_image = transform.postprocess(image)
+    return resized_image
 
-# Function to transform the image based on a prompt
 @spaces.GPU(enable_queue=True)
-def generate_image(image, prompt):
-    # Prepare the image for processing
-    image = image.convert("RGB")
-    recoloring_image = Image.fromarray(np.array(image)).convert('L').convert('RGB')
+def generate_image(input_image, prompt, controlnet_conditioning_scale):
+    # Always use a random seed for diversity in outputs
+    seed = np.random.randint(2147483647)
+    generator = torch.Generator("cuda").manual_seed(seed)  
+    
+    # Resize and prepare the image
+    input_image = resize_image(input_image)
+    grayscale_image = input_image.convert('L').convert('RGB')
     
-    # Define the negative prompt
-    negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
+    # Generate the image with fixed 30 steps
+    images = pipe(
+        prompt=prompt, 
+        image=grayscale_image, 
+        num_inference_steps=30, 
+        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
+        generator=generator,
+    ).images
     
-    # Generate the transformed image
-    results = pipe(prompt=prompt, negative_prompt=negative_prompt, image=recoloring_image, controlnet_conditioning_scale=1.0, height=1024, width=1024)
-    return results.images[0]
+    return images[0]
 
 # Gradio Interface
-description = """
-Anything to Anything, a workflow by Angrypenguinpng using the Bria Recolor ControlNet, check it out here: https://huggingface.co/briaai/BRIA-2.2-ControlNet-Recoloring
-"""
+description = "Anything to Anything. Transform anything to anything. Allow an adjuster for controlnet scale."
 
 with gr.Blocks() as demo:
     gr.Markdown("<h1><center>Image Transformation with Bria Recolor ControlNet</center></h1>")
     gr.Markdown(description)
-    with gr.Group():
-        with gr.Row():
-            image = gr.Image(label='Upload your image')
-            prompt = gr.Textbox(label='Enter your prompt', placeholder="A portrait of a beautiful and playful ethereal singer, golden designs, highly detailed, blurry background")
-            submit = gr.Button('Transform Image')
-    output_image = gr.Image(label='Transformed Image')
-
-    submit.click(fn=generate_image, inputs=[image, prompt], outputs=output_image)
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label='Upload your image', type="pil")
+            prompt = gr.Textbox(label='Enter your prompt')
+            controlnet_conditioning_scale = gr.Slider(label="ControlNet conditioning scale", minimum=0.1, maximum=2.0, value=1.0, step=0.05)
+            submit_button = gr.Button('Transform Image')
+        with gr.Column():
+            output_image = gr.Image(label='Transformed Image')
+
+    submit_button.click(fn=generate_image, inputs=[input_image, prompt, controlnet_conditioning_scale], outputs=output_image)
 
 demo.queue().launch()