app_py

app.py

@@ -7,7 +7,9 @@ from typing import Optional
 from diffusers import StableDiffusionPipeline, StableDiffusionControlNetPipeline
 from diffusers import ControlNetModel
 from peft import PeftModel, LoraConfig
-from PIL import Image
+from rembg import new_session, remove
+
+from PIL import Image as PILImage
 import cv2
 
 import torch
@@ -46,6 +48,7 @@ def infer(
     width,
     height,
     lscale=0.0,
+    remove_background=False,
     controlnet_enabled=False,
     controlnet_strength=0.0,
     controlnet_mode=None,
@@ -59,33 +62,32 @@ def infer(
     num_inference_steps : Optional[int] = 20,
     progress=gr.Progress(track_tqdm=True),
 ):
-    # if model_id != "CompVis/stable-diffusion-v1-4":
-    #    raise ValueError("The submitted model is not supported")
+
     generator = torch.Generator().manual_seed(seed)
 
     if controlnet_enabled:
         if not controlnet_image :
             raise ValueError("controlnet_enabled set to True, but controlnet_image not given")
         else:
-            controlnet_model = ControlNetModel.from_pretrained(CONTROL_MODE_MODEL.get(controlnet_mode))
+            controlnet_model = ControlNetModel.from_pretrained(CONTROL_MODE_MODEL.get(controlnet_mode),torch_dtype=torch_dtype)
         if model_id == "SD-v1-5 + Lora" :
             pipe=StableDiffusionControlNetPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5",controlnet=controlnet_model, torch_dtype=torch_dtype)
-            pipe.unet = PeftModel.from_pretrained("Emilichcka/diffusion_lora_funny_cat", "./unet", torch_dtype=torch_dtype)
-            pipe.text_encoder = PeftModel.from_pretrained("Emilichcka/diffusion_lora_funny_cat", "./text_encoder", torch_dtype=torch_dtype)
+            pipe.unet = PeftModel.from_pretrained(pipe.unet , "Emilichcka/diffusion_lora_funny_cat", subfolder="unet", torch_dtype=torch_dtype)
+            pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder,"Emilichcka/diffusion_lora_funny_cat", subfolder="text_encoder", torch_dtype=torch_dtype)
 
         else:
             pipe=StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=controlnet_model, torch_dtype=torch_dtype)
     else:
         if model_id == "SD-v1-5 + Lora" :
             pipe=StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5",torch_dtype=torch_dtype)
-            pipe.unet = PeftModel.from_pretrained("Emilichcka/diffusion_lora_funny_cat", "./unet", torch_dtype=torch_dtype)
-            pipe.text_encoder = PeftModel.from_pretrained("Emilichcka/diffusion_lora_funny_cat", "./text_encoder", torch_dtype=torch_dtype)
+            pipe.unet = PeftModel.from_pretrained(pipe.unet , "Emilichcka/diffusion_lora_funny_cat", subfolder="unet", torch_dtype=torch_dtype)
+            pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder,"Emilichcka/diffusion_lora_funny_cat", subfolder="text_encoder", torch_dtype=torch_dtype)
         else:
             pipe=StableDiffusionPipeline.from_pretrained(model_id)
 
     if ip_adapter_enabled:
             ip_adapter_scale = float(ip_adapter_scale)
-            pipe.load_ip_adapter("h94/IP-Adapter",subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
+            pipe.load_ip_adapter("h94/IP-Adapter",subfolder="models", weight_name="ip-adapter-plus_sd15.bin", torch_dtype=torch_dtype)
             pipe.set_ip_adapter_scale(ip_adapter_scale)
 
     if controlnet_image!= None:
@@ -97,26 +99,27 @@ def infer(
       controlnet_image = cv2.Canny(controlnet_image, low_threshold, high_threshold)
       controlnet_image = controlnet_image[:, :, None]
       controlnet_image = np.concatenate([controlnet_image, controlnet_image, controlnet_image], axis=2)
-      controlnet_image = Image.fromarray(controlnet_image)
+      controlnet_image = PILImage.fromarray(controlnet_image)
 
     pipe = pipe.to(device)
 
-    try:
-      image = pipe(
-          prompt=prompt,
-          image=controlnet_image,
-          negative_prompt=negative_prompt,
-          guidance_scale=guidance_scale,
-          num_inference_steps=num_inference_steps,
-          width=width,
-          height=height,
-          generator=generator,
-          ross_attention_kwargs={"scale": float(lscale)},
-          controlnet_conditioning_scale=controlnet_strength,
-          ip_adapter_image=ip_adapter_image,
-      ).images[0]
-    except Exception as e:
-            raise gr.Error(f"Ошибка при генерации изображения: {e}")
+
+    image = pipe(
+        prompt=prompt,
+        image=controlnet_image,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+        ross_attention_kwargs={"scale": float(lscale)},
+        controlnet_conditioning_scale=controlnet_strength,
+        ip_adapter_image=ip_adapter_image,
+    ).images[0]
+
+    if remove_background:
+      image = remove(image)
 
     return image, seed
 
@@ -138,6 +141,7 @@ default_model_id_choice = [
     "stable-diffusion-v1-5/stable-diffusion-v1-5",
     "CompVis/stable-diffusion-v1-4",
     "SD-v1-5 + Lora",
+    "nota-ai/bk-sdm-small",
 
 ]
 
@@ -155,7 +159,7 @@ with gr.Blocks(css=css) as demo:
             model_id = gr.Dropdown(
             label="Model Selection",
             choices=default_model_id_choice,
-            value="CompVis/stable-diffusion-v1-4",
+            value="SD-v1-5 + Lora",
             )
 
             seed = gr.Slider(
@@ -180,6 +184,7 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Row():
+            remove_background = gr.Checkbox(label="Remove Background", value=False)
             controlnet_enabled = gr.Checkbox(label="Enable ControlNet", value=False)
             ip_adapter_enabled = gr.Checkbox(label="Enable IP-Adapter", value=False)
 
@@ -218,7 +223,6 @@ with gr.Blocks(css=css) as demo:
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
                 label="Negative prompt",
-                value="deformed, ugly,low res, worst quality, low quality",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
             )
@@ -238,7 +242,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=512,  # Replace with defaults that work for your model
                 )
 
                 height = gr.Slider(
@@ -246,7 +250,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    value=512,  # Replace with defaults that work for your model
                 )
 
             with gr.Row():
@@ -255,7 +259,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=7.0,  # Replace with defaults that work for your model
+                    value=10.0,  # Replace with defaults that work for your model
                 )
 
                 num_inference_steps = gr.Slider(
@@ -263,7 +267,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=20,  # Replace with defaults that work for your model
+                    value=30,  # Replace with defaults that work for your model
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])
@@ -277,6 +281,7 @@ with gr.Blocks(css=css) as demo:
             width,
             height,
             lora_scale,
+            remove_background,
             controlnet_enabled,
             controlNet_strength,
             controlNet_mode,