Update app.py

app.py

@@ -4,7 +4,7 @@ import diffusers
 from diffusers.models import AutoencoderKL
 
 vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
-pipeline = diffusers.DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V5.1_noVAE", vae=vae).to("cuda")
+
 
 def read_content(file_path: str) -> str:
     """read the content of target file
@@ -15,7 +15,18 @@ def read_content(file_path: str) -> str:
     return content
 
 
-def predict(prompt, negative_prompt, guidance_scale, num_inference_steps, scheduler, lora, lora_weight):
+def predict(prompt, negative_prompt, guidance_scale, num_inference_steps,model, scheduler, lora, lora_weight):
+    pipeline = diffusers.DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V6.0_B1_noVAE", vae=vae).to("cuda")
+    pipeline.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
+    if model == "Realistic_V5.1":
+        pipeline = diffusers.DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V5.1_noVAE", vae=vae).to("cuda")
+    if model == "Realistic_V5.0":
+        pipeline = diffusers.DiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V5.0_noVAE", vae=vae).to("cuda")
+        pipeline.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
+    if model == "EpicRealism":
+        pipeline = diffusers.DiffusionPipeline.from_pretrained("emilianJR/epiCRealism", vae=vae).to("cuda")
+        pipeline.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
+    
     scheduler_class_name = scheduler.split("-")[0]
     add_kwargs = {}
     if len(scheduler.split("-")) > 1:
@@ -23,13 +34,17 @@ def predict(prompt, negative_prompt, guidance_scale, num_inference_steps, schedu
     if len(scheduler.split("-")) > 2:
         add_kwargs["algorithm_type"] = "sde-dpmsolver++"
     scheduler = getattr(diffusers, scheduler_class_name)
+    
     pipeline.scheduler = scheduler.from_pretrained("emilianJR/epiCRealism", subfolder="scheduler", **add_kwargs)
+    
     if lora == "add_detail":
         lora = "profaker/add_detail_lora"
     if lora == "nursing_job":
         lora = "profaker/Nursing_job_lora"
     if lora == "nsfw_POV":
         lora = "profaker/NSFW_POV_lora"
+    if lora == "nayanthara":
+        lora = "profaker/Naya_lora"
     if lora == "None":
         images = pipeline(
         prompt=prompt,
@@ -45,6 +60,7 @@ def predict(prompt, negative_prompt, guidance_scale, num_inference_steps, schedu
         return images
         
     pipeline.load_lora_weights(lora)
+    
     images = pipeline(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -52,6 +68,7 @@ def predict(prompt, negative_prompt, guidance_scale, num_inference_steps, schedu
         guidance_scale=guidance_scale,
         cross_attention_kwargs={"scale": lora_weight}
     ).images[0]
+    
     print("Prompt", prompt)
     print("Negative", negative_prompt)
     print("Steps", num_inference_steps)
@@ -103,6 +120,9 @@ with image_blocks as demo:
                 with gr.Row(equal_height=True):
                     negative_prompt = gr.Textbox(label="negative_prompt", placeholder="Your negative prompt",
                                                  info="what you don't want to see in the image")
+                with gr.Row(equal_height=True):
+                    models = ['Realistic_V6.0','Realistic_V5.1','Realistic_V5.0','EpicRealism']
+                    model = gr.Dropdown(label="Models",choices=models,value="Realistic_V6.0")
                 with gr.Row(equal_height=True):
                     schedulers = ["DEISMultistepScheduler", "HeunDiscreteScheduler", "EulerDiscreteScheduler",
                                   "DPMSolverMultistepScheduler", "DPMSolverMultistepScheduler-Karras",
@@ -110,18 +130,18 @@ with image_blocks as demo:
                     scheduler = gr.Dropdown(label="Schedulers", choices=schedulers,
                                             value="DPMSolverMultistepScheduler-Karras")
                 with gr.Row(equal_height=True):
-                    lora = ['None','add_detail', 'nursing_job', 'nsfw_POV']
-                    lora = gr.Dropdown(label='Lora', choices=lora, value="None")
-                    lora_weight = [-1, -0.5, 0, 0.5, 1]
-                    lora_weight = gr.Dropdown(label="Lora Weights", choices=lora_weight, value=0.5)
+                    loras = ['None','add_detail', 'nursing_job', 'nsfw_POV','nayanthara']
+                    lora = gr.Dropdown(label='Lora', choices=loras, value="None")
+                    lora_weights = [-1, -0.5, 0, 0.5, 1]
+                    lora_weight = gr.Dropdown(label="Lora Weights", choices=lora_weights, value=0.5)
             with gr.Row(equal_height=True):
                 btn = gr.Button("Generate", elem_id="run_button")
 
         with gr.Column():
             image_out = gr.Image(label="Output", elem_id="output-img", height=512, width=512)
-    btn.click(fn=predict, inputs=[prompt, negative_prompt, guidance_scale, steps, scheduler, lora, lora_weight],
+    btn.click(fn=predict, inputs=[prompt, negative_prompt, guidance_scale, steps, model,scheduler, lora, lora_weight],
               outputs=[image_out], api_name='run')
-    prompt.submit(fn=predict, inputs=[prompt, negative_prompt, guidance_scale, steps, scheduler, lora, lora_weight],
+    prompt.submit(fn=predict, inputs=[prompt, negative_prompt, guidance_scale, steps, model,scheduler, lora, lora_weight],
                   outputs=[image_out])
 
 image_blocks.queue(max_size=25, api_open=True).launch(show_api=True)