Update app.py

app.py

@@ -6,32 +6,24 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 repo = "fluently/Fluently-XL-Final"
+pipe = StableDiffusionXLPipeline.from_pretrained(repo, torch_dtype=torch.float16, vae=vae)
+pipe.load_lora_weights("KingNish/Better-Image-XL-Lora", weight_name="example-03.safetensors", adapter_name="lora")
+pipe.set_adapters("lora")
+pipe.to("cuda")
 
-pipe_best = StableDiffusionXLPipeline.from_pretrained(repo, torch_dtype=torch.float16, vae=vae)
-pipe_best.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
-pipe_best.load_lora_weights("KingNish/Better-Image-XL-Lora", weight_name="example-03.safetensors", adapter_name="lora")
-pipe_best.set_adapters(["lora","dalle"], adapter_weights=[1.5, 0.5])
-pipe_best.to("cuda")
-
-pipe_3D = StableDiffusionXLPipeline.from_pretrained(repo, torch_dtype=torch.float16, vae=vae)
-pipe_3D.load_lora_weights("artificialguybr/3DRedmond-V1", weight_name="3DRedmond-3DRenderStyle-3DRenderAF.safetensors", adapter_name="3D")
-pipe_3D.set_adapters(["3D"])
-pipe_3D.to("cuda")
-
-pipe_logo = StableDiffusionXLPipeline.from_pretrained(repo, torch_dtype=torch.float16, vae=vae)
-pipe_logo.load_lora_weights("artificialguybr/LogoRedmond-LogoLoraForSDXL", weight_name="LogoRedmond_LogoRedAF.safetensors", adapter_name="logo")
-pipe_logo.set_adapters(["logo"])
-pipe_logo.to("cuda")
+refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+refiner.to("cuda")
 
 help_text = """
 To optimize image results:
@@ -66,7 +58,7 @@ pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_
 pipe_edit.to("cuda")
 
 # Generator
-@spaces.GPU(duration=45, queue=False)
+@spaces.GPU(duration=30, queue=False)
 def king(type ,
         input_image ,
         instruction: str ,
@@ -77,7 +69,7 @@ def king(type ,
         image_cfg_scale: float = 1.7,
         width: int = 1024,
         height: int = 1024,
-        style="BEST",
+        guidance_scale: float = 6,
         use_resolution_binning: bool = True,
         progress=gr.Progress(track_tqdm=True),
     ):
@@ -99,14 +91,23 @@ def king(type ,
         if randomize_seed:
             seed = random.randint(0, 99999)
         generator = torch.Generator().manual_seed(seed)
-        if style=="3D":
-            instruction = f"3DRenderAF, 3D Render, {instruction}"
-            image = pipe_3D( prompt = instruction, guidance_scale = 5, num_inference_steps = steps, width = width, height = height, generator = generator).images[0]         
-        elif style=="Logo":
-            instruction = f"LogoRedAF, {instruction}"
-            image = pipe_logo( prompt = instruction, guidance_scale = 5, num_inference_steps = steps, width = width, height = height, generator = generator).images[0]         
-        else:
-            image = pipe_best( prompt = instruction, guidance_scale = 5, num_inference_steps = steps, width = width, height = height, generator = generator).images[0]         
+        image = pipe(
+            prompt = instruction,
+            guidance_scale =6, 
+            num_inference_steps = steps, 
+            width = width, 
+            height = height,
+            generator = generator
+            output_type="latent",
+        ).images         
+
+        refine = refiner(
+            prompt=instruction,
+            guidance_scale=6,
+            num_inference_steps=25,
+            image=image,
+            generator=generator,
+        ).images[0]        
         return seed, image
 
 client = InferenceClient()
@@ -176,15 +177,10 @@ with gr.Blocks(css=css) as demo:
             type = gr.Dropdown(["Image Generation","Image Editing"], label="Task", value="Image Generation",interactive=True, info="AI will select option based on your query, but if it selects wrong, please choose correct one.")
         with gr.Column(scale=1):
             generate_button = gr.Button("Generate")
-    with gr.Row():
-        style = gr.Radio(choices=["BEST","3D","Logo"],label="Style", value="BEST", interactive=True)
+
     with gr.Row():
         input_image = gr.Image(label="Image", type="pil", interactive=True)
 
-    with gr.Row():
-        width = gr.Number(value=1024, step=16,label="Width", interactive=True)
-        height = gr.Number(value=1024, step=16,label="Height", interactive=True)
-    
     with gr.Row():
         text_cfg_scale = gr.Number(value=7.3, step=0.1, label="Text CFG", interactive=True)
         image_cfg_scale = gr.Number(value=1.7, step=0.1,label="Image CFG", interactive=True)
@@ -207,7 +203,9 @@ with gr.Blocks(css=css) as demo:
     )
 
     gr.Markdown(help_text)
+
     instruction.change(fn=response, inputs=[instruction,input_image], outputs=type, queue=False)
+
     input_image.upload(fn=response, inputs=[instruction,input_image], outputs=type, queue=False)
     
     gr.on(triggers=[
@@ -223,9 +221,6 @@ with gr.Blocks(css=css) as demo:
                 seed,
                 text_cfg_scale,
                 image_cfg_scale,
-                width,
-                height,
-                style
             ],
             outputs=[seed, input_image],
         )