Update app.py

app.py

@@ -14,7 +14,7 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Load the diffusion pipeline
 pipe = StableDiffusionXLPipeline.from_single_file(
-    "https://huggingface.co/kayfahaarukku/AkashicPulse-v1.0/resolve/main/AkashicPulse-v1.0-ft-ft.safetensors",
+    "https://huggingface.co/kayfahaarukku/AkashicPulse-v1.0/blob/main/AkashicPulse-v1.0-ft-ft.safetensors",  # Fixed URL
     torch_dtype=torch.float16,
     custom_pipeline="lpw_stable_diffusion_xl",
     use_safetensors=True,
@@ -25,41 +25,46 @@ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.conf
 # Function to generate an image
 @spaces.GPU  # 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
-    
-    if randomize_seed:
-        seed = random.randint(0, 99999999)
-    if use_defaults:
-        prompt = f"{prompt}, best quality, amazing quality, very aesthetic"
-        negative_prompt = f"nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], {negative_prompt}"
-    generator = torch.manual_seed(seed)
-    
-    def callback(step, timestep, latents):
-        progress(step / num_inference_steps)
-        return
-    
-    width, height = map(int, resolution.split('x'))
-    image = pipe(
-        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]
+    try:
+        pipe.to('cuda')  # Move the model to GPU when the function is called
+        
+        if randomize_seed:
+            seed = random.randint(0, 99999999)
+        if use_defaults:
+            prompt = f"{prompt}, best quality, amazing quality, very aesthetic"
+            negative_prompt = f"nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], {negative_prompt}"
+        generator = torch.manual_seed(seed)
+        
+        def callback(step, timestep, latents):
+            progress(step / num_inference_steps)
+            return
+        
+        width, height = map(int, resolution.split('x'))
+        image = pipe(
+            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]
 
-    torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
 
-    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}"
+        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 image, seed, metadata_text
+    except Exception as e:
+        return None, seed, f"Error during generation: {str(e)}"
 
 # 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)
+    if image is None:
+        return gr.update(value=None), seed, gr.update(value=metadata_text)
     return image, seed, gr.update(value=metadata_text)
 
 def reset_inputs():
@@ -68,8 +73,10 @@ def reset_inputs():
 with gr.Blocks(title="irAsu 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
     gr.HTML(
         "<h1>irAsu 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."
-        )
+        "<p>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.</p>"
+    )
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
@@ -93,11 +100,7 @@ with gr.Blocks(title="irAsu 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
         with gr.Column():
             output_image = gr.Image(type="pil", label="Generated Image")
             with gr.Accordion("Parameters", open=False):
-                gr.Markdown(
-                    """
-                    This parameter is compatible with Stable Diffusion WebUI's parameter importer.
-                    """
-                )
+                gr.Markdown("This parameter is compatible with Stable Diffusion WebUI's parameter importer.")
                 metadata_textbox = gr.Textbox(lines=6, label="Image Parameters", interactive=False, max_lines=6)
             gr.Markdown(
                 """
@@ -116,7 +119,8 @@ with gr.Blocks(title="irAsu 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
     generate_button.click(
         interface_fn,
         inputs=[
-            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
+            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]
     )
@@ -125,7 +129,8 @@ with gr.Blocks(title="irAsu 1.0 Demo", theme="NoCrypt/[email protected]") as demo:
         reset_inputs,
         inputs=[],
         outputs=[
-            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input, metadata_textbox
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input,
+            guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input, metadata_textbox
         ]
     )