Create app.py

app.py

@@ -0,0 +1,76 @@
+import os
+import torch
+from torch import autocast
+from diffusers import StableDiffusionPipeline
+import gradio as gr
+
+# Model configuration
+model_path = "path/to/your/checkpoint.safetensors"  # Update this with your checkpoint path
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the model
+pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    safety_checker=None
+)
+pipe.to(device)
+
+# If you have a custom checkpoint, load it
+if os.path.exists(model_path):
+    pipe.unet.load_state_dict(torch.load(model_path))
+
+def generate_image(prompt, negative_prompt, num_steps, guidance_scale, width, height, seed):
+    """
+    Generate an image using Stable Diffusion
+    """
+    if seed == -1:
+        seed = int.from_bytes(os.urandom(2), "big")
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    with autocast(device):
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance_scale,
+            width=width,
+            height=height,
+            generator=generator
+        ).images[0]
+    
+    return image, seed
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Stable Diffusion 1.5 Custom Model")
+    
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
+            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here...")
+            
+            with gr.Row():
+                num_steps = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Number of Steps")
+                guidance_scale = gr.Slider(minimum=1, maximum=20, value=7.5, step=0.5, label="Guidance Scale")
+            
+            with gr.Row():
+                width = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Width")
+                height = gr.Slider(minimum=256, maximum=1024, value=512, step=64, label="Height")
+            
+            seed = gr.Number(label="Seed (-1 for random)", value=-1)
+            generate_btn = gr.Button("Generate Image")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+            used_seed = gr.Number(label="Used Seed")
+
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[prompt, negative_prompt, num_steps, guidance_scale, width, height, seed],
+        outputs=[output_image, used_seed]
+    )
+
+# Launch app locally
+if __name__ == "__main__":
+    demo.launch()