Model X demo

app.py

@@ -9,32 +9,55 @@ import tqdm
 # Enable TQDM progress tracking
 tqdm.monitor_interval = 0
 
-# Load the diffusion pipeline
-pipe = StableDiffusionXLPipeline.from_pretrained(
+# Load the diffusion pipelines
+pipe1 = StableDiffusionXLPipeline.from_pretrained(
     "kayfahaarukku/UrangDiffusion-1.0", 
-    torch_dtype=torch.float32, 
+    torch_dtype=torch.float16, 
     custom_pipeline="lpw_stable_diffusion_xl",
 )
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe1.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe1.scheduler.config)
 
-# Function to generate an image
-@spaces.GPU(duration=90)  # Adjust the duration as needed
-def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
-    pipe.to('cuda')  # Move the model to GPU when the function is called
+pipe2 = StableDiffusionXLPipeline.from_pretrained(
+    "kayfahaarukku/model-x", 
+    torch_dtype=torch.float16, 
+    custom_pipeline="lpw_stable_diffusion_xl",
+)
+pipe2.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe2.scheduler.config)
+
+# Function to generate images from both models
+@spaces.GPU
+def generate_comparison(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
+    pipe1.to('cuda')
+    pipe2.to('cuda')
     
     if randomize_seed:
         seed = random.randint(0, 99999999)
     if use_defaults:
-        prompt = f"{prompt}, masterpiece, best quality"
-        negative_prompt = f"nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, {negative_prompt}"
+        prompt = f"{prompt}, masterpiece, best quality, amazing quality, very aesthetic"
+        negative_prompt = f"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, displeasing, {negative_prompt}"
     generator = torch.manual_seed(seed)
     
     def callback(step, timestep, latents):
-        progress(step / num_inference_steps)
+        progress(step / (2 * num_inference_steps))
         return
     
     width, height = map(int, resolution.split('x'))
-    image = pipe(
+    
+    # Generate image with UrangDiffusion-1.0
+    image1 = pipe1(
+        prompt, 
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height, 
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        callback=callback,
+        callback_steps=1
+    ).images[0]
+
+    # Generate image with model-x
+    image2 = pipe2(
         prompt, 
         negative_prompt=negative_prompt,
         width=width,
@@ -50,20 +73,20 @@ def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_s
 
     metadata_text = f"{prompt}\nNegative prompt: {negative_prompt}\nSteps: {num_inference_steps}, Sampler: Euler a, Size: {width}x{height}, Seed: {seed}, CFG scale: {guidance_scale}"
 
-    return image, seed, metadata_text
+    return image1, image2, seed, metadata_text
 
 # Define Gradio interface
 def interface_fn(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
-    image, seed, metadata_text = generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
-    return image, seed, gr.update(value=metadata_text)
+    image1, image2, seed, metadata_text = generate_comparison(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
+    return image1, image2, seed, gr.update(value=metadata_text)
 
 def reset_inputs():
-    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='832x1216'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=True), gr.update(value='')
+    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='896x1152'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=True), gr.update(value='')
 
-with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
+with gr.Blocks(title="Model X Comparison Demo", theme="NoCrypt/[email protected]") as demo:
     gr.HTML(
-        "<h1>UrangDiffusion 1.0 Demo</h1>"
-        "This demo is intended to showcase what the model is capable of and is not intended to be the main generation platform. Results produced with Diffusers are not the best, and it's highly recommended for you to get the model running inside Stable Diffusion WebUI or ComfyUI."
+        "<h1>UrangDiffusion 1.0 vs Model X Comparison Demo</h1>"
+        "This demo showcases a comparison between UrangDiffusion 1.0 and Model X."
         )
     with gr.Row():
         with gr.Column():
@@ -76,17 +99,19 @@ with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/[email protected]") as d
                     "1344x768", "768x1344", "1536x640", "640x1536"
                 ],
                 label="Resolution",
-                value="832x1216"
+                value="896x1152"
             )
             guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7)
             num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
             seed_input = gr.Slider(minimum=0, maximum=99999999, step=1, label="Seed", value=0, interactive=True)
             randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
-            generate_button = gr.Button("Generate")
+            generate_button = gr.Button("Generate Comparison")
             reset_button = gr.Button("Reset")
 
         with gr.Column():
-            output_image = gr.Image(type="pil", label="Generated Image")
+            with gr.Row():
+                output_image1 = gr.Image(type="pil", label="UrangDiffusion 1.0")
+                output_image2 = gr.Image(type="pil", label="Model X")
             with gr.Accordion("Parameters", open=False):
                 gr.Markdown(
                     """
@@ -112,7 +137,7 @@ with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/[email protected]") as d
         inputs=[
             prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
         ],
-        outputs=[output_image, seed_input, metadata_textbox]
+        outputs=[output_image1, output_image2, seed_input, metadata_textbox]
     )
     
     reset_button.click(