Update app.py

app.py

@@ -1,4 +1,4 @@
-import torch
+rt torch
 import spaces
 from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
 from transformers import AutoFeatureExtractor
@@ -11,12 +11,11 @@ import gradio as gr
 import cv2
 
 base_model_paths = {
-    "RealisticVisionV4": "SG161222/Realistic_Vision_V4.0_noVAE",
-    "RealisticVisionV6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
+    "Realistic_Vision_V4.0_noVAE": "SG161222/Realistic_Vision_V4.0_noVAE",
+    "Realistic_Vision_V6.0_B1_noVAE": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
     "Deliberate": "Yntec/Deliberate",
-    "DeliberateV2": "Yntec/Deliberate2",
-    "Dreamshaper8": "Lykon/dreamshaper-8",
-    "EpicRealism": "emilianJR/epiCRealism"
+    "Deliberate2": "Yntec/Deliberate2",
+    "pony-diffusion": "AstraliteHeart/pony-diffusion"
 }
 
 vae_model_path = "stabilityai/sd-vae-ft-mse"
@@ -61,7 +60,7 @@ app.prepare(ctx_id=0, det_size=(640, 640))
 cv2.setNumThreads(1)
 
 @spaces.GPU(enable_queue=True)
-def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, base_model, progress=gr.Progress(track_tqdm=True)):
+def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, base_model, num_inference_steps, guidance_scale, width, height, progress=gr.Progress(track_tqdm=True)):
     global ip_model, ip_model_plus
     base_model_path = base_model_paths[base_model]
     pipe = load_model(base_model_path)
@@ -87,13 +86,13 @@ def generate_image(images, prompt, negative_prompt, preserve_face_structure, fac
         print("Generating normal")
         image = ip_model.generate(
             prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
-            scale=likeness_strength, width=512, height=512, num_inference_steps=30
+            scale=likeness_strength, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
         )
     else:
         print("Generating plus")
         image = ip_model_plus.generate(
             prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
-            scale=likeness_strength, face_image=face_image, shortcut=True, s_scale=face_strength, width=512, height=512, num_inference_steps=30
+            scale=likeness_strength, face_image=face_image, shortcut=True, s_scale=face_strength, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
         )
     print(image)
     return image
@@ -138,7 +137,11 @@ with gr.Blocks(css=css) as demo:
                 preserve = gr.Checkbox(label="Preserve Face Structure", info="Higher quality, less versatility (the face structure of your first photo will be preserved). Unchecking this will use the v1 model.", value=True)
                 face_strength = gr.Slider(label="Face Structure strength", info="Only applied if preserve face structure is checked", value=1.3, step=0.1, minimum=0, maximum=3)
                 likeness_strength = gr.Slider(label="Face Embed strength", value=1.0, step=0.1, minimum=0, maximum=5)
-                nfaa_negative_prompts = gr.Textbox(label="Appended Negative Prompts", info="Negative prompts to steer generations towards safe for all audiences outputs", value="naked, bikini, skimpy, scanty, bare skin, lingerie, swimsuit, exposed, see-through")    
+                nfaa_negative_prompts = gr.Textbox(label="Appended Negative Prompts", info="Negative prompts to steer generations towards safe for all audiences outputs", value="naked, bikini, skimpy, scanty, bare skin, lingerie, swimsuit, exposed, see-through")
+                num_inference_steps = gr.Slider(label="Number of Inference Steps", value=30, step=1, minimum=10, maximum=100)
+                guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, step=0.1, minimum=1, maximum=20)
+                width = gr.Slider(label="Width", value=512, step=64, minimum=256, maximum=1024)
+                height = gr.Slider(label="Height", value=512, step=64, minimum=256, maximum=1024)
         with gr.Column():
             gallery = gr.Gallery(label="Generated Images")
         style.change(fn=change_style,
@@ -147,7 +150,7 @@ with gr.Blocks(css=css) as demo:
         files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
         remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
         submit.click(fn=generate_image,
-                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts, base_model],
+                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts, base_model, num_inference_steps, guidance_scale, width, height],
                     outputs=gallery)
     
     gr.Markdown("")