File size: 1,707 Bytes
18274c1
453ed2e
 
 
1a833ba
453ed2e
a29e3ba
00f6a78
9ad92f4
453ed2e
9ad92f4
4984c7e
 
be85eb8
 
453ed2e
e56af76
b31f6c0
be85eb8
453ed2e
00f6a78
 
a29e3ba
4984c7e
be85eb8
766763f
324658b
 
 
 
 
 
 
766763f
9ad92f4
00f6a78
453ed2e
00f6a78
324658b
 
 
4984c7e
453ed2e
7391723
324658b
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import spaces
import torch
import gradio as gr
from PIL import Image
import random
from diffusers import (
    DiffusionPipeline,
    AutoencoderKL,
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionControlNetImg2ImgPipeline,
    DPMSolverMultistepScheduler,
    EulerDiscreteScheduler
)
import tempfile
import time
import os

BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"

# Initialize both pipelines
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16)

# Commenting out safety checker initialization
# 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("cuda")
#     feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")

main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
    BASE_MODEL,
    controlnet=controlnet,
    vae=vae,
    # Remove safety checker and feature extractor from pipeline components
    # safety_checker=safety_checker,
    # feature_extractor=feature_extractor,
    torch_dtype=torch.float16,
).to("cuda")

def inference(control_image: Image.Image, prompt: str, negative_prompt: str, guidance_scale: float = 8.0, controlnet_conditioning_scale: float = 1, upscaler_strength: float = 0.5, seed: int = -1):
    # Inference logic remains unchanged
    ...