Add application file

README.md

@@ -2,7 +2,7 @@
 title: IllusionDiffusion
 emoji: 🔥
 colorFrom: green
-colorTo: pink
+colorTo: red
 sdk: gradio
 sdk_version: 4.36.1
 app_file: app.py

app.py

@@ -27,16 +27,14 @@ from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionS
 
 BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"
 
-# Initialize both pipelines
 vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float32)
 controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float32)
 
-# Initialize the safety checker conditionally
 SAFETY_CHECKER_ENABLED = os.environ.get("SAFETY_CHECKER", "0") == "1"
 safety_checker = None
 feature_extractor = None
 if SAFETY_CHECKER_ENABLED:
-    safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to("cpu")
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
     feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
 main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
@@ -46,11 +44,10 @@ main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
     safety_checker=safety_checker,
     feature_extractor=feature_extractor,
     torch_dtype=torch.float32,
-).to("cpu")
+)
 
 image_pipe = StableDiffusionControlNetImg2ImgPipeline(**main_pipe.components)
 
-# Sampler map
 SAMPLER_MAP = {
     "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
     "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
@@ -58,18 +55,13 @@ SAMPLER_MAP = {
 
 def center_crop_resize(img, output_size=(512, 512)):
     width, height = img.size
-
-    # Calculate dimensions to crop to the center
     new_dimension = min(width, height)
     left = (width - new_dimension)/2
     top = (height - new_dimension)/2
     right = (width + new_dimension)/2
     bottom = (height + new_dimension)/2
-
-    # Crop and resize
     img = img.crop((left, top, right, bottom))
     img = img.resize(output_size)
-
     return img
 
 def common_upscale(samples, width, height, upscale_method, crop=False):
@@ -87,7 +79,6 @@ def common_upscale(samples, width, height, upscale_method, crop=False):
             s = samples[:,:,y:old_height-y,x:old_width-x]
         else:
             s = samples
-
         return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
 
 def upscale(samples, upscale_method, scale_by):
@@ -111,7 +102,6 @@ def convert_to_base64(pil_image):
         image.save(temp_file.name)
     return temp_file.name
 
-# Inference function
 def inference(
     control_image: Image.Image,
     prompt: str,
@@ -136,7 +126,7 @@ def inference(
 
     main_pipe.scheduler = SAMPLER_MAP[sampler](main_pipe.scheduler.config)
     my_seed = random.randint(0, 2**32 - 1) if seed == -1 else seed
-    generator = torch.Generator(device="cpu").manual_seed(my_seed)
+    generator = torch.Generator().manual_seed(my_seed)
     
     out = main_pipe(
         prompt=prompt,
@@ -169,7 +159,6 @@ def inference(
     end_time_formatted = time.strftime("%H:%M:%S", end_time_struct)
     print(f"Inference ended at {end_time_formatted}, taking {end_time-start_time}s")
 
-    # Save image + metadata
     user_history.save_image(
         label=prompt,
         image=out_image["images"][0],